/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.arrow.flight.sql.example;

import static com.google.common.base.Strings.emptyToNull;
import static com.google.protobuf.Any.pack;
import static com.google.protobuf.ByteString.copyFrom;
import static java.lang.String.format;
import static java.util.Collections.singletonList;
import static java.util.Objects.isNull;
import static java.util.UUID.randomUUID;
import static java.util.stream.IntStream.range;
import static org.apache.arrow.adapter.jdbc.JdbcToArrow.sqlToArrowVectorIterator;
import static org.apache.arrow.adapter.jdbc.JdbcToArrowUtils.jdbcToArrowSchema;
import static org.apache.arrow.flight.sql.impl.FlightSql.*;
import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCrossReference;
import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetDbSchemas;
import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetExportedKeys;
import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetImportedKeys;
import static org.apache.arrow.flight.sql.impl.FlightSql.DoPutUpdateResult;
import static org.apache.arrow.flight.sql.impl.FlightSql.TicketStatementQuery;
import static org.apache.arrow.util.Preconditions.checkState;
import static org.slf4j.LoggerFactory.getLogger;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.RemovalListener;
import com.google.common.cache.RemovalNotification;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.protobuf.ByteString;
import com.google.protobuf.Message;
import com.google.protobuf.ProtocolStringList;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.NoSuchFileException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.nio.file.StandardOpenOption;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.SQLSyntaxErrorException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Calendar;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
import java.util.Properties;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.apache.arrow.adapter.jdbc.ArrowVectorIterator;
import org.apache.arrow.adapter.jdbc.JdbcFieldInfo;
import org.apache.arrow.adapter.jdbc.JdbcParameterBinder;
import org.apache.arrow.adapter.jdbc.JdbcToArrowUtils;
import org.apache.arrow.flight.CallStatus;
import org.apache.arrow.flight.Criteria;
import org.apache.arrow.flight.FlightDescriptor;
import org.apache.arrow.flight.FlightEndpoint;
import org.apache.arrow.flight.FlightInfo;
import org.apache.arrow.flight.FlightServer;
import org.apache.arrow.flight.FlightStream;
import org.apache.arrow.flight.Location;
import org.apache.arrow.flight.PutResult;
import org.apache.arrow.flight.Result;
import org.apache.arrow.flight.SchemaResult;
import org.apache.arrow.flight.Ticket;
import org.apache.arrow.flight.sql.FlightSqlColumnMetadata;
import org.apache.arrow.flight.sql.FlightSqlProducer;
import org.apache.arrow.flight.sql.SqlInfoBuilder;
import org.apache.arrow.flight.sql.impl.FlightSql.ActionClosePreparedStatementRequest;
import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementRequest;
import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementResult;
import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCatalogs;
import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetPrimaryKeys;
import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetSqlInfo;
import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTableTypes;
import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTables;
import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementQuery;
import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementUpdate;
import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest;
import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDefinitionOptions;
import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDefinitionOptions.TableExistsOption;
import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDefinitionOptions.TableNotExistOption;
import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementQuery;
import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementUpdate;
import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedCaseSensitivity;
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.util.Preconditions;
import org.apache.arrow.vector.BitVector;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.UInt1Vector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorLoader;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.VectorUnloader;
import org.apache.arrow.vector.complex.ListVector;
import org.apache.arrow.vector.complex.impl.UnionListWriter;
import org.apache.arrow.vector.ipc.WriteChannel;
import org.apache.arrow.vector.ipc.message.MessageSerializer;
import org.apache.arrow.vector.types.Types.MinorType;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.arrow.vector.util.Text;
import org.apache.commons.dbcp2.ConnectionFactory;
import org.apache.commons.dbcp2.DriverManagerConnectionFactory;
import org.apache.commons.dbcp2.PoolableConnection;
import org.apache.commons.dbcp2.PoolableConnectionFactory;
import org.apache.commons.dbcp2.PoolingDataSource;
import org.apache.commons.pool2.ObjectPool;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.apache.commons.text.StringEscapeUtils;
import org.slf4j.Logger;

/**
 * Example {@link FlightSqlProducer} implementation showing an Apache Derby backed Flight SQL server
 * that generally supports all current features of Flight SQL.
 */
public class FlightSqlExample implements FlightSqlProducer, AutoCloseable {
  private static final Logger LOGGER = getLogger(FlightSqlExample.class);
  protected static final Calendar DEFAULT_CALENDAR = JdbcToArrowUtils.getUtcCalendar();
  public static final String DB_NAME = "derbyDB";
  private final String databaseUri;
  // ARROW-15315: Use ExecutorService to simulate an async scenario
  private final ExecutorService executorService = Executors.newFixedThreadPool(10);
  private final Location location;
  protected final PoolingDataSource<PoolableConnection> dataSource;
  protected final BufferAllocator rootAllocator = new RootAllocator();
  private final Cache<ByteString, StatementContext<PreparedStatement>>
      preparedStatementLoadingCache;
  private final Cache<ByteString, StatementContext<Statement>> statementLoadingCache;
  private final SqlInfoBuilder sqlInfoBuilder;

  public static void main(String[] args) throws Exception {
    Location location = Location.forGrpcInsecure("localhost", 55555);
    final FlightSqlExample example = new FlightSqlExample(location, DB_NAME);
    Location listenLocation = Location.forGrpcInsecure("0.0.0.0", 55555);
    try (final BufferAllocator allocator = new RootAllocator();
        final FlightServer server =
            FlightServer.builder(allocator, listenLocation, example).build()) {
      server.start();
      server.awaitTermination();
    }
  }

  public FlightSqlExample(final Location location, final String dbName) {
    // TODO Constructor should not be doing work.
    checkState(removeDerbyDatabaseIfExists(dbName), "Failed to clear Derby database!");
    checkState(populateDerbyDatabase(dbName), "Failed to populate Derby database!");
    databaseUri = "jdbc:derby:target/" + dbName;
    final ConnectionFactory connectionFactory =
        new DriverManagerConnectionFactory(databaseUri, new Properties());
    final PoolableConnectionFactory poolableConnectionFactory =
        new PoolableConnectionFactory(connectionFactory, null);
    final ObjectPool<PoolableConnection> connectionPool =
        new GenericObjectPool<>(poolableConnectionFactory);

    poolableConnectionFactory.setPool(connectionPool);
    // PoolingDataSource takes ownership of `connectionPool`
    dataSource = new PoolingDataSource<>(connectionPool);

    preparedStatementLoadingCache =
        CacheBuilder.newBuilder()
            .maximumSize(100)
            .expireAfterWrite(10, TimeUnit.MINUTES)
            .removalListener(new StatementRemovalListener<PreparedStatement>())
            .build();

    statementLoadingCache =
        CacheBuilder.newBuilder()
            .maximumSize(100)
            .expireAfterWrite(10, TimeUnit.MINUTES)
            .removalListener(new StatementRemovalListener<>())
            .build();

    this.location = location;

    sqlInfoBuilder = new SqlInfoBuilder();
    try (final Connection connection = dataSource.getConnection()) {
      final DatabaseMetaData metaData = connection.getMetaData();

      sqlInfoBuilder
          .withFlightSqlServerName(metaData.getDatabaseProductName())
          .withFlightSqlServerVersion(metaData.getDatabaseProductVersion())
          .withFlightSqlServerArrowVersion(metaData.getDriverVersion())
          .withFlightSqlServerReadOnly(metaData.isReadOnly())
          .withFlightSqlServerSql(true)
          .withFlightSqlServerSubstrait(false)
          .withFlightSqlServerTransaction(SqlSupportedTransaction.SQL_SUPPORTED_TRANSACTION_NONE)
          .withSqlIdentifierQuoteChar(metaData.getIdentifierQuoteString())
          .withSqlDdlCatalog(metaData.supportsCatalogsInDataManipulation())
          .withSqlDdlSchema(metaData.supportsSchemasInDataManipulation())
          .withSqlDdlTable(metaData.allTablesAreSelectable())
          .withSqlIdentifierCase(
              metaData.storesMixedCaseIdentifiers()
                  ? SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_CASE_INSENSITIVE
                  : metaData.storesUpperCaseIdentifiers()
                      ? SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_UPPERCASE
                      : metaData.storesLowerCaseIdentifiers()
                          ? SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_LOWERCASE
                          : SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_UNKNOWN)
          .withSqlQuotedIdentifierCase(
              metaData.storesMixedCaseQuotedIdentifiers()
                  ? SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_CASE_INSENSITIVE
                  : metaData.storesUpperCaseQuotedIdentifiers()
                      ? SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_UPPERCASE
                      : metaData.storesLowerCaseQuotedIdentifiers()
                          ? SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_LOWERCASE
                          : SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_UNKNOWN)
          .withSqlAllTablesAreSelectable(true)
          .withSqlNullOrdering(SqlNullOrdering.SQL_NULLS_SORTED_AT_END)
          .withSqlMaxColumnsInTable(42)
          .withFlightSqlServerBulkIngestion(true)
          .withFlightSqlServerBulkIngestionTransaction(false);
    } catch (SQLException e) {
      throw new RuntimeException(e);
    }
  }

  public static boolean removeDerbyDatabaseIfExists(final String dbName) {
    final Path path = Paths.get("target" + File.separator + dbName);

    try (final Stream<Path> walk = Files.walk(path)) {
      /*
       * Iterate over all paths to delete, mapping each path to the outcome of its own
       * deletion as a boolean representing whether each individual operation was
       * successful; then reduce all booleans into a single answer.
       * If for whatever reason the resulting `Stream<Boolean>` is empty, throw an `IOException`;
       * this not expected.
       */
      boolean unused =
          walk.sorted(Comparator.reverseOrder())
              .map(Path::toFile)
              .map(File::delete)
              .reduce(Boolean::logicalAnd)
              .orElseThrow(IOException::new);
    } catch (NoSuchFileException e) {
      /*
       * The only acceptable scenario for an `IOException` to be thrown here is if
       * an attempt to delete an non-existing file takes place -- which should be
       * alright, since they would be deleted anyway.
       */
      LOGGER.error(format("No existing Derby database to delete.: <%s>", e.getMessage()), e);
      return true;
    } catch (Exception e) {
      LOGGER.error(format("Failed attempt to clear DerbyDB.: <%s>", e.getMessage()), e);
      return false;
    }
    return true;
  }

  private static boolean populateDerbyDatabase(final String dbName) {
    try (final Connection connection =
            DriverManager.getConnection("jdbc:derby:target/" + dbName + ";create=true");
        Statement statement = connection.createStatement()) {

      dropTable(statement, "intTable");
      dropTable(statement, "foreignTable");
      statement.execute(
          "CREATE TABLE foreignTable ("
              + "id INT not null primary key GENERATED ALWAYS AS IDENTITY (START WITH 1, INCREMENT BY 1), "
              + "foreignName varchar(100), "
              + "value int)");
      statement.execute(
          "CREATE TABLE intTable ("
              + "id INT not null primary key GENERATED ALWAYS AS IDENTITY (START WITH 1, INCREMENT BY 1), "
              + "keyName varchar(100), "
              + "value int, "
              + "foreignId int references foreignTable(id))");
      statement.execute("INSERT INTO foreignTable (foreignName, value) VALUES ('keyOne', 1)");
      statement.execute("INSERT INTO foreignTable (foreignName, value) VALUES ('keyTwo', 0)");
      statement.execute("INSERT INTO foreignTable (foreignName, value) VALUES ('keyThree', -1)");
      statement.execute("INSERT INTO intTable (keyName, value, foreignId) VALUES ('one', 1, 1)");
      statement.execute("INSERT INTO intTable (keyName, value, foreignId) VALUES ('zero', 0, 1)");
      statement.execute(
          "INSERT INTO intTable (keyName, value, foreignId) VALUES ('negative one', -1, 1)");
    } catch (final SQLException e) {
      LOGGER.error(format("Failed attempt to populate DerbyDB: <%s>", e.getMessage()), e);
      return false;
    }
    return true;
  }

  private static void dropTable(final Statement statement, final String tableName)
      throws SQLException {
    try {
      statement.execute("DROP TABLE " + tableName);
    } catch (SQLException e) {
      // sql error code for "object does not exist"; which is fine, we're trying to delete the table
      // see https://db.apache.org/derby/docs/10.17/ref/rrefexcept71493.html
      if (!"42Y55".equals(e.getSQLState())) {
        throw e;
      }
    }
  }

  private static ArrowType getArrowTypeFromJdbcType(
      final int jdbcDataType, final int precision, final int scale) {
    try {
      return JdbcToArrowUtils.getArrowTypeFromJdbcType(
          new JdbcFieldInfo(jdbcDataType, precision, scale), DEFAULT_CALENDAR);
    } catch (UnsupportedOperationException ignored) {
      return ArrowType.Utf8.INSTANCE;
    }
  }

  private static void saveToVector(final Byte data, final UInt1Vector vector, final int index) {
    vectorConsumer(
        data,
        vector,
        fieldVector -> fieldVector.setNull(index),
        (theData, fieldVector) -> fieldVector.setSafe(index, theData));
  }

  private static void saveToVector(final Byte data, final BitVector vector, final int index) {
    vectorConsumer(
        data,
        vector,
        fieldVector -> fieldVector.setNull(index),
        (theData, fieldVector) -> fieldVector.setSafe(index, theData));
  }

  private static void saveToVector(final String data, final VarCharVector vector, final int index) {
    preconditionCheckSaveToVector(vector, index);
    vectorConsumer(
        data,
        vector,
        fieldVector -> fieldVector.setNull(index),
        (theData, fieldVector) -> fieldVector.setSafe(index, new Text(theData)));
  }

  private static void saveToVector(final Integer data, final IntVector vector, final int index) {
    preconditionCheckSaveToVector(vector, index);
    vectorConsumer(
        data,
        vector,
        fieldVector -> fieldVector.setNull(index),
        (theData, fieldVector) -> fieldVector.setSafe(index, theData));
  }

  private static void saveToVector(
      final byte[] data, final VarBinaryVector vector, final int index) {
    preconditionCheckSaveToVector(vector, index);
    vectorConsumer(
        data,
        vector,
        fieldVector -> fieldVector.setNull(index),
        (theData, fieldVector) -> fieldVector.setSafe(index, theData));
  }

  private static void preconditionCheckSaveToVector(final FieldVector vector, final int index) {
    Objects.requireNonNull(vector, "vector cannot be null.");
    checkState(index >= 0, "Index must be a positive number!");
  }

  private static <T, V extends FieldVector> void vectorConsumer(
      final T data,
      final V vector,
      final Consumer<V> consumerIfNullable,
      final BiConsumer<T, V> defaultConsumer) {
    if (isNull(data)) {
      consumerIfNullable.accept(vector);
      return;
    }
    defaultConsumer.accept(data, vector);
  }

  private static VectorSchemaRoot getSchemasRoot(
      final ResultSet data, final BufferAllocator allocator) throws SQLException {
    final VarCharVector catalogs = new VarCharVector("catalog_name", allocator);
    final VarCharVector schemas =
        new VarCharVector(
            "db_schema_name", FieldType.notNullable(MinorType.VARCHAR.getType()), allocator);
    final List<FieldVector> vectors = ImmutableList.of(catalogs, schemas);
    vectors.forEach(FieldVector::allocateNew);
    final Map<FieldVector, String> vectorToColumnName =
        ImmutableMap.of(
            catalogs, "TABLE_CATALOG",
            schemas, "TABLE_SCHEM");
    saveToVectors(vectorToColumnName, data);
    final int rows =
        vectors.stream()
            .map(FieldVector::getValueCount)
            .findAny()
            .orElseThrow(IllegalStateException::new);
    vectors.forEach(vector -> vector.setValueCount(rows));
    return new VectorSchemaRoot(vectors);
  }

  private static <T extends FieldVector> int saveToVectors(
      final Map<T, String> vectorToColumnName, final ResultSet data, boolean emptyToNull)
      throws SQLException {
    Predicate<ResultSet> alwaysTrue = (resultSet) -> true;
    return saveToVectors(vectorToColumnName, data, emptyToNull, alwaysTrue);
  }

  @SuppressWarnings("StringSplitter")
  private static <T extends FieldVector> int saveToVectors(
      final Map<T, String> vectorToColumnName,
      final ResultSet data,
      boolean emptyToNull,
      Predicate<ResultSet> resultSetPredicate)
      throws SQLException {
    Objects.requireNonNull(vectorToColumnName, "vectorToColumnName cannot be null.");
    Objects.requireNonNull(data, "data cannot be null.");
    final Set<Entry<T, String>> entrySet = vectorToColumnName.entrySet();
    int rows = 0;

    while (data.next()) {
      if (!resultSetPredicate.test(data)) {
        continue;
      }
      for (final Entry<T, String> vectorToColumn : entrySet) {
        final T vector = vectorToColumn.getKey();
        final String columnName = vectorToColumn.getValue();
        if (vector instanceof VarCharVector) {
          String thisData = data.getString(columnName);
          saveToVector(
              emptyToNull ? emptyToNull(thisData) : thisData, (VarCharVector) vector, rows);
        } else if (vector instanceof IntVector) {
          final int intValue = data.getInt(columnName);
          saveToVector(data.wasNull() ? null : intValue, (IntVector) vector, rows);
        } else if (vector instanceof UInt1Vector) {
          final byte byteValue = data.getByte(columnName);
          saveToVector(data.wasNull() ? null : byteValue, (UInt1Vector) vector, rows);
        } else if (vector instanceof BitVector) {
          final byte byteValue = data.getByte(columnName);
          saveToVector(data.wasNull() ? null : byteValue, (BitVector) vector, rows);
        } else if (vector instanceof ListVector) {
          String createParamsValues = data.getString(columnName);

          UnionListWriter writer = ((ListVector) vector).getWriter();

          BufferAllocator allocator = vector.getAllocator();
          final ArrowBuf buf = allocator.buffer(1024);

          writer.setPosition(rows);
          writer.startList();

          if (createParamsValues != null) {
            String[] split = createParamsValues.split(",");

            range(0, split.length)
                .forEach(
                    i -> {
                      byte[] bytes = split[i].getBytes(StandardCharsets.UTF_8);
                      Preconditions.checkState(
                          bytes.length < 1024,
                          "The amount of bytes is greater than what the ArrowBuf supports");
                      buf.setBytes(0, bytes);
                      writer.varChar().writeVarChar(0, bytes.length, buf);
                    });
          }
          buf.close();
          writer.endList();
        } else {
          throw CallStatus.INVALID_ARGUMENT
              .withDescription("Provided vector not supported")
              .toRuntimeException();
        }
      }
      rows++;
    }
    for (final Entry<T, String> vectorToColumn : entrySet) {
      vectorToColumn.getKey().setValueCount(rows);
    }

    return rows;
  }

  private static <T extends FieldVector> void saveToVectors(
      final Map<T, String> vectorToColumnName, final ResultSet data) throws SQLException {
    saveToVectors(vectorToColumnName, data, false);
  }

  private static VectorSchemaRoot getTableTypesRoot(
      final ResultSet data, final BufferAllocator allocator) throws SQLException {
    return getRoot(data, allocator, "table_type", "TABLE_TYPE");
  }

  private static VectorSchemaRoot getCatalogsRoot(
      final ResultSet data, final BufferAllocator allocator) throws SQLException {
    return getRoot(data, allocator, "catalog_name", "TABLE_CATALOG");
  }

  private static VectorSchemaRoot getRoot(
      final ResultSet data,
      final BufferAllocator allocator,
      final String fieldVectorName,
      final String columnName)
      throws SQLException {
    final VarCharVector dataVector =
        new VarCharVector(
            fieldVectorName, FieldType.notNullable(MinorType.VARCHAR.getType()), allocator);
    saveToVectors(ImmutableMap.of(dataVector, columnName), data);
    final int rows = dataVector.getValueCount();
    dataVector.setValueCount(rows);
    return new VectorSchemaRoot(singletonList(dataVector));
  }

  private static VectorSchemaRoot getTypeInfoRoot(
      CommandGetXdbcTypeInfo request, ResultSet typeInfo, final BufferAllocator allocator)
      throws SQLException {
    Preconditions.checkNotNull(allocator, "BufferAllocator cannot be null.");

    VectorSchemaRoot root = VectorSchemaRoot.create(Schemas.GET_TYPE_INFO_SCHEMA, allocator);

    Map<FieldVector, String> mapper = new HashMap<>();
    mapper.put(root.getVector("type_name"), "TYPE_NAME");
    mapper.put(root.getVector("data_type"), "DATA_TYPE");
    mapper.put(root.getVector("column_size"), "PRECISION");
    mapper.put(root.getVector("literal_prefix"), "LITERAL_PREFIX");
    mapper.put(root.getVector("literal_suffix"), "LITERAL_SUFFIX");
    mapper.put(root.getVector("create_params"), "CREATE_PARAMS");
    mapper.put(root.getVector("nullable"), "NULLABLE");
    mapper.put(root.getVector("case_sensitive"), "CASE_SENSITIVE");
    mapper.put(root.getVector("searchable"), "SEARCHABLE");
    mapper.put(root.getVector("unsigned_attribute"), "UNSIGNED_ATTRIBUTE");
    mapper.put(root.getVector("fixed_prec_scale"), "FIXED_PREC_SCALE");
    mapper.put(root.getVector("auto_increment"), "AUTO_INCREMENT");
    mapper.put(root.getVector("local_type_name"), "LOCAL_TYPE_NAME");
    mapper.put(root.getVector("minimum_scale"), "MINIMUM_SCALE");
    mapper.put(root.getVector("maximum_scale"), "MAXIMUM_SCALE");
    mapper.put(root.getVector("sql_data_type"), "SQL_DATA_TYPE");
    mapper.put(root.getVector("datetime_subcode"), "SQL_DATETIME_SUB");
    mapper.put(root.getVector("num_prec_radix"), "NUM_PREC_RADIX");

    Predicate<ResultSet> predicate;
    if (request.hasDataType()) {
      predicate =
          (resultSet) -> {
            try {
              return resultSet.getInt("DATA_TYPE") == request.getDataType();
            } catch (SQLException e) {
              throw new RuntimeException(e);
            }
          };
    } else {
      predicate = resultSet -> true;
    }

    int rows = saveToVectors(mapper, typeInfo, true, predicate);

    root.setRowCount(rows);
    return root;
  }

  private static VectorSchemaRoot getTablesRoot(
      final DatabaseMetaData databaseMetaData,
      final BufferAllocator allocator,
      final boolean includeSchema,
      final String catalog,
      final String schemaFilterPattern,
      final String tableFilterPattern,
      final String... tableTypes)
      throws SQLException, IOException {
    /*
     * TODO Fix DerbyDB inconsistency if possible.
     * During the early development of this prototype, an inconsistency has been found in the database
     * used for this demonstration; as DerbyDB does not operate with the concept of catalogs, fetching
     * the catalog name for a given table from `DatabaseMetadata#getColumns` and `DatabaseMetadata#getSchemas`
     * returns null, as expected. However, the inconsistency lies in the fact that accessing the same
     * information -- that is, the catalog name for a given table -- from `DatabaseMetadata#getSchemas`
     * returns an empty String.The temporary workaround for this was making sure we convert the empty Strings
     * to null using `com.google.common.base.Strings#emptyToNull`.
     */
    Objects.requireNonNull(allocator, "BufferAllocator cannot be null.");
    final VarCharVector catalogNameVector = new VarCharVector("catalog_name", allocator);
    final VarCharVector schemaNameVector = new VarCharVector("db_schema_name", allocator);
    final VarCharVector tableNameVector =
        new VarCharVector(
            "table_name", FieldType.notNullable(MinorType.VARCHAR.getType()), allocator);
    final VarCharVector tableTypeVector =
        new VarCharVector(
            "table_type", FieldType.notNullable(MinorType.VARCHAR.getType()), allocator);

    final List<FieldVector> vectors = new ArrayList<>(4);
    vectors.add(catalogNameVector);
    vectors.add(schemaNameVector);
    vectors.add(tableNameVector);
    vectors.add(tableTypeVector);

    vectors.forEach(FieldVector::allocateNew);

    final Map<FieldVector, String> vectorToColumnName =
        ImmutableMap.of(
            catalogNameVector, "TABLE_CAT",
            schemaNameVector, "TABLE_SCHEM",
            tableNameVector, "TABLE_NAME",
            tableTypeVector, "TABLE_TYPE");

    try (final ResultSet data =
        Objects.requireNonNull(
                databaseMetaData,
                format("%s cannot be null.", databaseMetaData.getClass().getName()))
            .getTables(catalog, schemaFilterPattern, tableFilterPattern, tableTypes)) {

      saveToVectors(vectorToColumnName, data, true);
      final int rows =
          vectors.stream()
              .map(FieldVector::getValueCount)
              .findAny()
              .orElseThrow(IllegalStateException::new);
      vectors.forEach(vector -> vector.setValueCount(rows));

      if (includeSchema) {
        final VarBinaryVector tableSchemaVector =
            new VarBinaryVector(
                "table_schema", FieldType.notNullable(MinorType.VARBINARY.getType()), allocator);
        tableSchemaVector.allocateNew(rows);

        try (final ResultSet columnsData =
            databaseMetaData.getColumns(catalog, schemaFilterPattern, tableFilterPattern, null)) {
          final Map<String, List<Field>> tableToFields = new HashMap<>();

          while (columnsData.next()) {
            final String catalogName = columnsData.getString("TABLE_CAT");
            final String schemaName = columnsData.getString("TABLE_SCHEM");
            final String tableName = columnsData.getString("TABLE_NAME");
            final String typeName = columnsData.getString("TYPE_NAME");
            final String fieldName = columnsData.getString("COLUMN_NAME");
            final int dataType = columnsData.getInt("DATA_TYPE");
            final boolean isNullable =
                columnsData.getInt("NULLABLE") != DatabaseMetaData.columnNoNulls;
            final int precision = columnsData.getInt("COLUMN_SIZE");
            final int scale = columnsData.getInt("DECIMAL_DIGITS");
            boolean isAutoIncrement =
                Objects.equals(columnsData.getString("IS_AUTOINCREMENT"), "YES");

            final List<Field> fields =
                tableToFields.computeIfAbsent(tableName, tableName_ -> new ArrayList<>());

            final FlightSqlColumnMetadata columnMetadata =
                new FlightSqlColumnMetadata.Builder()
                    .catalogName(catalogName)
                    .schemaName(schemaName)
                    .tableName(tableName)
                    .typeName(typeName)
                    .precision(precision)
                    .scale(scale)
                    .isAutoIncrement(isAutoIncrement)
                    .build();

            final Field field =
                new Field(
                    fieldName,
                    new FieldType(
                        isNullable,
                        getArrowTypeFromJdbcType(dataType, precision, scale),
                        null,
                        columnMetadata.getMetadataMap()),
                    null);
            fields.add(field);
          }

          for (int index = 0; index < rows; index++) {
            final String tableName = tableNameVector.getObject(index).toString();
            final Schema schema = new Schema(tableToFields.get(tableName));
            saveToVector(
                copyFrom(serializeMetadata(schema)).toByteArray(), tableSchemaVector, index);
          }
        }

        tableSchemaVector.setValueCount(rows);
        vectors.add(tableSchemaVector);
      }
    }

    return new VectorSchemaRoot(vectors);
  }

  private static ByteBuffer serializeMetadata(final Schema schema) {
    final ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
    try {
      MessageSerializer.serialize(new WriteChannel(Channels.newChannel(outputStream)), schema);

      return ByteBuffer.wrap(outputStream.toByteArray());
    } catch (final IOException e) {
      throw new RuntimeException("Failed to serialize schema", e);
    }
  }

  private static String getRootAsCSVNoHeader(final VectorSchemaRoot root) {
    StringBuilder sb = new StringBuilder();
    Schema schema = root.getSchema();
    int rowCount = root.getRowCount();
    List<FieldVector> fieldVectors = root.getFieldVectors();

    List<Object> row = new ArrayList<>(schema.getFields().size());
    for (int i = 0; i < rowCount; i++) {
      if (i > 0) {
        sb.append("\n");
      }
      row.clear();
      for (FieldVector v : fieldVectors) {
        row.add(v.getObject(i));
      }
      printRowAsCSV(sb, row);
    }
    return sb.toString();
  }

  private static void printRowAsCSV(StringBuilder sb, List<Object> values) {
    sb.append(
        values.stream()
            .map(v -> isNull(v) ? "" : v.toString())
            .map(StringEscapeUtils::escapeCsv)
            .collect(Collectors.joining(",")));
  }

  @Override
  public void getStreamPreparedStatement(
      final CommandPreparedStatementQuery command,
      final CallContext context,
      final ServerStreamListener listener) {
    final ByteString handle = command.getPreparedStatementHandle();
    StatementContext<PreparedStatement> statementContext =
        preparedStatementLoadingCache.getIfPresent(handle);
    Objects.requireNonNull(statementContext);
    final PreparedStatement statement = statementContext.getStatement();
    try (final ResultSet resultSet = statement.executeQuery()) {
      final Schema schema = jdbcToArrowSchema(resultSet.getMetaData(), DEFAULT_CALENDAR);
      try (final VectorSchemaRoot vectorSchemaRoot =
          VectorSchemaRoot.create(schema, rootAllocator)) {
        final VectorLoader loader = new VectorLoader(vectorSchemaRoot);
        listener.start(vectorSchemaRoot);

        final ArrowVectorIterator iterator = sqlToArrowVectorIterator(resultSet, rootAllocator);
        while (iterator.hasNext()) {
          final VectorSchemaRoot batch = iterator.next();
          if (batch.getRowCount() == 0) {
            break;
          }
          final VectorUnloader unloader = new VectorUnloader(batch);
          loader.load(unloader.getRecordBatch());
          listener.putNext();
          vectorSchemaRoot.clear();
        }

        listener.putNext();
      }
    } catch (final SQLException | IOException e) {
      LOGGER.error(format("Failed to getStreamPreparedStatement: <%s>.", e.getMessage()), e);
      listener.error(
          CallStatus.INTERNAL
              .withDescription("Failed to prepare statement: " + e)
              .toRuntimeException());
    } finally {
      listener.completed();
    }
  }

  @Override
  public void closePreparedStatement(
      final ActionClosePreparedStatementRequest request,
      final CallContext context,
      final StreamListener<Result> listener) {
    // Running on another thread
    Future<?> unused =
        executorService.submit(
            () -> {
              try {
                preparedStatementLoadingCache.invalidate(request.getPreparedStatementHandle());
              } catch (final Exception e) {
                listener.onError(e);
                return;
              }
              listener.onCompleted();
            });
  }

  @Override
  public FlightInfo getFlightInfoStatement(
      final CommandStatementQuery request,
      final CallContext context,
      final FlightDescriptor descriptor) {
    ByteString handle = copyFrom(randomUUID().toString().getBytes(StandardCharsets.UTF_8));

    try {
      // Ownership of the connection will be passed to the context. Do NOT close!
      final Connection connection = dataSource.getConnection();
      final Statement statement =
          connection.createStatement(ResultSet.TYPE_SCROLL_INSENSITIVE, ResultSet.CONCUR_READ_ONLY);
      final String query = request.getQuery();
      final StatementContext<Statement> statementContext = new StatementContext<>(statement, query);

      statementLoadingCache.put(handle, statementContext);
      final ResultSet resultSet = statement.executeQuery(query);

      TicketStatementQuery ticket =
          TicketStatementQuery.newBuilder().setStatementHandle(handle).build();
      return getFlightInfoForSchema(
          ticket, descriptor, jdbcToArrowSchema(resultSet.getMetaData(), DEFAULT_CALENDAR));
    } catch (final SQLException e) {
      LOGGER.error(
          format("There was a problem executing the prepared statement: <%s>.", e.getMessage()), e);
      throw CallStatus.INTERNAL.withCause(e).toRuntimeException();
    }
  }

  @Override
  public FlightInfo getFlightInfoPreparedStatement(
      final CommandPreparedStatementQuery command,
      final CallContext context,
      final FlightDescriptor descriptor) {
    final ByteString preparedStatementHandle = command.getPreparedStatementHandle();
    StatementContext<PreparedStatement> statementContext =
        preparedStatementLoadingCache.getIfPresent(preparedStatementHandle);
    try {
      assert statementContext != null;
      PreparedStatement statement = statementContext.getStatement();

      ResultSetMetaData metaData = statement.getMetaData();
      return getFlightInfoForSchema(
          command, descriptor, jdbcToArrowSchema(metaData, DEFAULT_CALENDAR));
    } catch (final SQLException e) {
      LOGGER.error(
          format("There was a problem executing the prepared statement: <%s>.", e.getMessage()), e);
      throw CallStatus.INTERNAL.withCause(e).toRuntimeException();
    }
  }

  @Override
  public SchemaResult getSchemaStatement(
      final CommandStatementQuery command,
      final CallContext context,
      final FlightDescriptor descriptor) {
    throw CallStatus.UNIMPLEMENTED.toRuntimeException();
  }

  @Override
  public void close() throws Exception {
    try {
      preparedStatementLoadingCache.cleanUp();
    } catch (Throwable t) {
      LOGGER.error(format("Failed to close resources: <%s>", t.getMessage()), t);
    }

    AutoCloseables.close(dataSource, rootAllocator);
  }

  @Override
  public void listFlights(
      CallContext context, Criteria criteria, StreamListener<FlightInfo> listener) {
    // TODO - build example implementation
    throw CallStatus.UNIMPLEMENTED.toRuntimeException();
  }

  @Override
  public void createPreparedStatement(
      final ActionCreatePreparedStatementRequest request,
      final CallContext context,
      final StreamListener<Result> listener) {
    // Running on another thread
    Future<?> unused =
        executorService.submit(
            () -> {
              try {
                final ByteString preparedStatementHandle =
                    copyFrom(request.getQuery().getBytes(StandardCharsets.UTF_8));
                // Ownership of the connection will be passed to the context. Do NOT close!
                final Connection connection = dataSource.getConnection();
                final PreparedStatement preparedStatement =
                    connection.prepareStatement(
                        request.getQuery(),
                        ResultSet.TYPE_SCROLL_INSENSITIVE,
                        ResultSet.CONCUR_READ_ONLY);
                final StatementContext<PreparedStatement> preparedStatementContext =
                    new StatementContext<>(preparedStatement, request.getQuery());

                preparedStatementLoadingCache.put(
                    preparedStatementHandle, preparedStatementContext);

                final Schema parameterSchema =
                    jdbcToArrowSchema(preparedStatement.getParameterMetaData(), DEFAULT_CALENDAR);

                final ResultSetMetaData metaData = preparedStatement.getMetaData();
                final ByteString bytes =
                    isNull(metaData)
                        ? ByteString.EMPTY
                        : ByteString.copyFrom(
                            serializeMetadata(jdbcToArrowSchema(metaData, DEFAULT_CALENDAR)));
                final ActionCreatePreparedStatementResult result =
                    ActionCreatePreparedStatementResult.newBuilder()
                        .setDatasetSchema(bytes)
                        .setParameterSchema(copyFrom(serializeMetadata(parameterSchema)))
                        .setPreparedStatementHandle(preparedStatementHandle)
                        .build();
                listener.onNext(new Result(pack(result).toByteArray()));
              } catch (final SQLException e) {
                listener.onError(
                    CallStatus.INTERNAL
                        .withDescription("Failed to create prepared statement: " + e)
                        .toRuntimeException());
                return;
              } catch (final Throwable t) {
                listener.onError(
                    CallStatus.INTERNAL
                        .withDescription("Unknown error: " + t)
                        .toRuntimeException());
                return;
              }
              listener.onCompleted();
            });
  }

  @Override
  public void doExchange(CallContext context, FlightStream reader, ServerStreamListener writer) {
    // TODO - build example implementation
    throw CallStatus.UNIMPLEMENTED.toRuntimeException();
  }

  @Override
  public Runnable acceptPutStatement(
      CommandStatementUpdate command,
      CallContext context,
      FlightStream flightStream,
      StreamListener<PutResult> ackStream) {
    final String query = command.getQuery();

    return () -> {
      try (final Connection connection = dataSource.getConnection();
          final Statement statement = connection.createStatement()) {
        final int result = statement.executeUpdate(query);

        final DoPutUpdateResult build =
            DoPutUpdateResult.newBuilder().setRecordCount(result).build();

        try (final ArrowBuf buffer = rootAllocator.buffer(build.getSerializedSize())) {
          buffer.writeBytes(build.toByteArray());
          ackStream.onNext(PutResult.metadata(buffer));
          ackStream.onCompleted();
        }
      } catch (SQLSyntaxErrorException e) {
        ackStream.onError(
            CallStatus.INVALID_ARGUMENT
                .withDescription("Failed to execute statement (invalid syntax): " + e)
                .toRuntimeException());
      } catch (SQLException e) {
        ackStream.onError(
            CallStatus.INTERNAL
                .withDescription("Failed to execute statement: " + e)
                .toRuntimeException());
      }
    };
  }

  @Override
  public Runnable acceptPutStatementBulkIngest(
      CommandStatementIngest command,
      CallContext context,
      FlightStream flightStream,
      StreamListener<PutResult> ackStream) {

    final String schema = command.hasSchema() ? command.getSchema() : null;
    final String table = command.getTable();
    final boolean temporary = command.getTemporary();
    final boolean transactionId = command.hasTransactionId();
    final TableDefinitionOptions tableDefinitionOptions =
        command.hasTableDefinitionOptions() ? command.getTableDefinitionOptions() : null;

    return () -> {
      TableExistsOption ifExists = TableExistsOption.TABLE_EXISTS_OPTION_APPEND;
      if (temporary) {
        ackStream.onError(
            CallStatus.UNIMPLEMENTED
                .withDescription("Bulk ingestion using temporary tables is not supported")
                .toRuntimeException());
      } else if (transactionId) {
        ackStream.onError(
            CallStatus.UNIMPLEMENTED
                .withDescription(
                    "Bulk ingestion automatically happens in a transaction. Specifying explicit transaction is not supported.")
                .toRuntimeException());
      } else if (isNull(tableDefinitionOptions)) {
        ackStream.onError(
            CallStatus.INVALID_ARGUMENT
                .withDescription("TableDefinitionOptions not provided.")
                .toRuntimeException());
      } else {
        TableNotExistOption ifNotExist = tableDefinitionOptions.getIfNotExist();
        ifExists = tableDefinitionOptions.getIfExists();

        if (!TableNotExistOption.TABLE_NOT_EXIST_OPTION_FAIL.equals(ifNotExist)) {
          ackStream.onError(
              CallStatus.UNIMPLEMENTED
                  .withDescription(
                      "Only supported option is TABLE_NOT_EXIST_OPTION_FAIL for TableNotExistsOption.")
                  .toRuntimeException());
        } else if (TableExistsOption.TABLE_EXISTS_OPTION_UNSPECIFIED.equals(ifExists)) {
          ackStream.onError(
              CallStatus.INVALID_ARGUMENT
                  .withDescription("TableExistsOption must be specified")
                  .toRuntimeException());
        } else if (TableExistsOption.TABLE_EXISTS_OPTION_FAIL.equals(ifExists)) {
          ackStream.onError(
              CallStatus.UNIMPLEMENTED
                  .withDescription("TABLE_EXISTS_OPTION_FAIL is not supported.")
                  .toRuntimeException());
        }
      }

      Path tempFile = null;
      try {
        tempFile = Files.createTempFile(null, null);

        VectorSchemaRoot root = null;
        int counter = 0;
        while (flightStream.next()) {
          if (counter > 0) {
            Files.writeString(tempFile, "\n", StandardCharsets.UTF_8, StandardOpenOption.APPEND);
          }
          counter += 1;
          root = flightStream.getRoot();
          Files.writeString(
              tempFile,
              getRootAsCSVNoHeader(root),
              StandardCharsets.UTF_8,
              StandardOpenOption.APPEND);
        }

        if (counter > 0) {
          Files.writeString(tempFile, "\n", StandardCharsets.UTF_8, StandardOpenOption.APPEND);
        }

        if (!isNull(root)) {
          String header =
              root.getSchema().getFields().stream()
                  .map(Field::getName)
                  .collect(Collectors.joining(","));

          try (final Connection connection = dataSource.getConnection();
              final PreparedStatement preparedStatement =
                  connection.prepareStatement(
                      "CALL SYSCS_UTIL.SYSCS_IMPORT_DATA (?,?,?,null,?,?,?,?,?)")) {

            preparedStatement.setString(1, schema);
            preparedStatement.setString(2, table);
            preparedStatement.setString(3, header);
            preparedStatement.setString(4, tempFile.toString());
            preparedStatement.setString(5, ",");
            preparedStatement.setString(6, "\"");
            preparedStatement.setString(7, "UTF-8");
            preparedStatement.setInt(
                8, TableExistsOption.TABLE_EXISTS_OPTION_REPLACE.equals(ifExists) ? 1 : 0);
            preparedStatement.execute();

            final DoPutUpdateResult build =
                DoPutUpdateResult.newBuilder().setRecordCount(-1).build();

            try (final ArrowBuf buffer = rootAllocator.buffer(build.getSerializedSize())) {
              buffer.writeBytes(build.toByteArray());
              ackStream.onNext(PutResult.metadata(buffer));
              ackStream.onCompleted();
            }
          } catch (SQLException e) {
            ackStream.onError(
                CallStatus.INTERNAL
                    .withDescription("Failed to execute bulk ingest: " + e)
                    .toRuntimeException());
          }
        }
      } catch (IOException e) {
        ackStream.onError(
            CallStatus.INTERNAL
                .withDescription("Failed to create temp file for bulk loading: " + e)
                .toRuntimeException());
      } finally {
        if (!isNull(tempFile)) {
          try {
            Files.delete(tempFile);
          } catch (IOException e) {
            //
          }
        }
      }
    };
  }

  @Override
  public Runnable acceptPutPreparedStatementUpdate(
      CommandPreparedStatementUpdate command,
      CallContext context,
      FlightStream flightStream,
      StreamListener<PutResult> ackStream) {
    final StatementContext<PreparedStatement> statement =
        preparedStatementLoadingCache.getIfPresent(command.getPreparedStatementHandle());

    return () -> {
      if (statement == null) {
        ackStream.onError(
            CallStatus.NOT_FOUND
                .withDescription("Prepared statement does not exist")
                .toRuntimeException());
        return;
      }
      try {
        final PreparedStatement preparedStatement = statement.getStatement();

        while (flightStream.next()) {
          final VectorSchemaRoot root = flightStream.getRoot();

          final int rowCount = root.getRowCount();
          final int recordCount;

          if (rowCount == 0) {
            preparedStatement.execute();
            recordCount = preparedStatement.getUpdateCount();
          } else {
            final JdbcParameterBinder binder =
                JdbcParameterBinder.builder(preparedStatement, root).bindAll().build();
            while (binder.next()) {
              preparedStatement.addBatch();
            }
            final int[] recordCounts = preparedStatement.executeBatch();
            recordCount = Arrays.stream(recordCounts).sum();
          }

          final DoPutUpdateResult build =
              DoPutUpdateResult.newBuilder().setRecordCount(recordCount).build();

          try (final ArrowBuf buffer = rootAllocator.buffer(build.getSerializedSize())) {
            buffer.writeBytes(build.toByteArray());
            ackStream.onNext(PutResult.metadata(buffer));
          }
        }
      } catch (SQLException e) {
        ackStream.onError(
            CallStatus.INTERNAL
                .withDescription("Failed to execute update: " + e)
                .toRuntimeException());
        return;
      }
      ackStream.onCompleted();
    };
  }

  @Override
  public Runnable acceptPutPreparedStatementQuery(
      CommandPreparedStatementQuery command,
      CallContext context,
      FlightStream flightStream,
      StreamListener<PutResult> ackStream) {
    final StatementContext<PreparedStatement> statementContext =
        preparedStatementLoadingCache.getIfPresent(command.getPreparedStatementHandle());

    return () -> {
      assert statementContext != null;
      PreparedStatement preparedStatement = statementContext.getStatement();

      try {
        while (flightStream.next()) {
          final VectorSchemaRoot root = flightStream.getRoot();
          final JdbcParameterBinder binder =
              JdbcParameterBinder.builder(preparedStatement, root).bindAll().build();
          while (binder.next()) {
            // Do not execute() - will be done in a getStream call
          }
        }

      } catch (SQLException e) {
        ackStream.onError(
            CallStatus.INTERNAL
                .withDescription("Failed to bind parameters: " + e.getMessage())
                .withCause(e)
                .toRuntimeException());
        return;
      }

      ackStream.onCompleted();
    };
  }

  @Override
  public FlightInfo getFlightInfoSqlInfo(
      final CommandGetSqlInfo request,
      final CallContext context,
      final FlightDescriptor descriptor) {
    return getFlightInfoForSchema(request, descriptor, Schemas.GET_SQL_INFO_SCHEMA);
  }

  @Override
  public void getStreamSqlInfo(
      final CommandGetSqlInfo command,
      final CallContext context,
      final ServerStreamListener listener) {
    this.sqlInfoBuilder.send(command.getInfoList(), listener);
  }

  @Override
  public FlightInfo getFlightInfoTypeInfo(
      CommandGetXdbcTypeInfo request, CallContext context, FlightDescriptor descriptor) {
    return getFlightInfoForSchema(request, descriptor, Schemas.GET_TYPE_INFO_SCHEMA);
  }

  @Override
  public void getStreamTypeInfo(
      CommandGetXdbcTypeInfo request, CallContext context, ServerStreamListener listener) {
    try (final Connection connection = dataSource.getConnection();
        final ResultSet typeInfo = connection.getMetaData().getTypeInfo();
        final VectorSchemaRoot vectorSchemaRoot =
            getTypeInfoRoot(request, typeInfo, rootAllocator)) {
      listener.start(vectorSchemaRoot);
      listener.putNext();
    } catch (SQLException e) {
      LOGGER.error(format("Failed to getStreamCatalogs: <%s>.", e.getMessage()), e);
      listener.error(e);
    } finally {
      listener.completed();
    }
  }

  @Override
  public FlightInfo getFlightInfoCatalogs(
      final CommandGetCatalogs request,
      final CallContext context,
      final FlightDescriptor descriptor) {
    return getFlightInfoForSchema(request, descriptor, Schemas.GET_CATALOGS_SCHEMA);
  }

  @Override
  public void getStreamCatalogs(final CallContext context, final ServerStreamListener listener) {
    try (final Connection connection = dataSource.getConnection();
        final ResultSet catalogs = connection.getMetaData().getCatalogs();
        final VectorSchemaRoot vectorSchemaRoot = getCatalogsRoot(catalogs, rootAllocator)) {
      listener.start(vectorSchemaRoot);
      listener.putNext();
    } catch (SQLException e) {
      LOGGER.error(format("Failed to getStreamCatalogs: <%s>.", e.getMessage()), e);
      listener.error(e);
    } finally {
      listener.completed();
    }
  }

  @Override
  public FlightInfo getFlightInfoSchemas(
      final CommandGetDbSchemas request,
      final CallContext context,
      final FlightDescriptor descriptor) {
    return getFlightInfoForSchema(request, descriptor, Schemas.GET_SCHEMAS_SCHEMA);
  }

  @Override
  public void getStreamSchemas(
      final CommandGetDbSchemas command,
      final CallContext context,
      final ServerStreamListener listener) {
    final String catalog = command.hasCatalog() ? command.getCatalog() : null;
    final String schemaFilterPattern =
        command.hasDbSchemaFilterPattern() ? command.getDbSchemaFilterPattern() : null;
    try (final Connection connection = dataSource.getConnection();
        final ResultSet schemas =
            connection.getMetaData().getSchemas(catalog, schemaFilterPattern);
        final VectorSchemaRoot vectorSchemaRoot = getSchemasRoot(schemas, rootAllocator)) {
      listener.start(vectorSchemaRoot);
      listener.putNext();
    } catch (SQLException e) {
      LOGGER.error(format("Failed to getStreamSchemas: <%s>.", e.getMessage()), e);
      listener.error(e);
    } finally {
      listener.completed();
    }
  }

  @Override
  public FlightInfo getFlightInfoTables(
      final CommandGetTables request,
      final CallContext context,
      final FlightDescriptor descriptor) {
    Schema schemaToUse = Schemas.GET_TABLES_SCHEMA;
    if (!request.getIncludeSchema()) {
      schemaToUse = Schemas.GET_TABLES_SCHEMA_NO_SCHEMA;
    }
    return getFlightInfoForSchema(request, descriptor, schemaToUse);
  }

  @Override
  public void getStreamTables(
      final CommandGetTables command,
      final CallContext context,
      final ServerStreamListener listener) {
    final String catalog = command.hasCatalog() ? command.getCatalog() : null;
    final String schemaFilterPattern =
        command.hasDbSchemaFilterPattern() ? command.getDbSchemaFilterPattern() : null;
    final String tableFilterPattern =
        command.hasTableNameFilterPattern() ? command.getTableNameFilterPattern() : null;

    final ProtocolStringList protocolStringList = command.getTableTypesList();
    final int protocolSize = protocolStringList.size();
    final String[] tableTypes =
        protocolSize == 0 ? null : protocolStringList.toArray(new String[protocolSize]);

    try (final Connection connection = DriverManager.getConnection(databaseUri);
        final VectorSchemaRoot vectorSchemaRoot =
            getTablesRoot(
                connection.getMetaData(),
                rootAllocator,
                command.getIncludeSchema(),
                catalog,
                schemaFilterPattern,
                tableFilterPattern,
                tableTypes)) {
      listener.start(vectorSchemaRoot);
      listener.putNext();
    } catch (SQLException | IOException e) {
      LOGGER.error(format("Failed to getStreamTables: <%s>.", e.getMessage()), e);
      listener.error(e);
    } finally {
      listener.completed();
    }
  }

  @Override
  public FlightInfo getFlightInfoTableTypes(
      final CommandGetTableTypes request,
      final CallContext context,
      final FlightDescriptor descriptor) {
    return getFlightInfoForSchema(request, descriptor, Schemas.GET_TABLE_TYPES_SCHEMA);
  }

  @Override
  public void getStreamTableTypes(final CallContext context, final ServerStreamListener listener) {
    try (final Connection connection = dataSource.getConnection();
        final ResultSet tableTypes = connection.getMetaData().getTableTypes();
        final VectorSchemaRoot vectorSchemaRoot = getTableTypesRoot(tableTypes, rootAllocator)) {
      listener.start(vectorSchemaRoot);
      listener.putNext();
    } catch (SQLException e) {
      LOGGER.error(format("Failed to getStreamTableTypes: <%s>.", e.getMessage()), e);
      listener.error(e);
    } finally {
      listener.completed();
    }
  }

  @Override
  public FlightInfo getFlightInfoPrimaryKeys(
      final CommandGetPrimaryKeys request,
      final CallContext context,
      final FlightDescriptor descriptor) {
    return getFlightInfoForSchema(request, descriptor, Schemas.GET_PRIMARY_KEYS_SCHEMA);
  }

  @Override
  public void getStreamPrimaryKeys(
      final CommandGetPrimaryKeys command,
      final CallContext context,
      final ServerStreamListener listener) {

    final String catalog = command.hasCatalog() ? command.getCatalog() : null;
    final String schema = command.hasDbSchema() ? command.getDbSchema() : null;
    final String table = command.getTable();

    try (Connection connection = DriverManager.getConnection(databaseUri)) {
      final ResultSet primaryKeys = connection.getMetaData().getPrimaryKeys(catalog, schema, table);

      final VarCharVector catalogNameVector = new VarCharVector("catalog_name", rootAllocator);
      final VarCharVector schemaNameVector = new VarCharVector("db_schema_name", rootAllocator);
      final VarCharVector tableNameVector = new VarCharVector("table_name", rootAllocator);
      final VarCharVector columnNameVector = new VarCharVector("column_name", rootAllocator);
      final IntVector keySequenceVector = new IntVector("key_sequence", rootAllocator);
      final VarCharVector keyNameVector = new VarCharVector("key_name", rootAllocator);

      final List<FieldVector> vectors =
          new ArrayList<>(
              ImmutableList.of(
                  catalogNameVector,
                  schemaNameVector,
                  tableNameVector,
                  columnNameVector,
                  keySequenceVector,
                  keyNameVector));
      vectors.forEach(FieldVector::allocateNew);

      int rows = 0;
      for (; primaryKeys.next(); rows++) {
        saveToVector(primaryKeys.getString("TABLE_CAT"), catalogNameVector, rows);
        saveToVector(primaryKeys.getString("TABLE_SCHEM"), schemaNameVector, rows);
        saveToVector(primaryKeys.getString("TABLE_NAME"), tableNameVector, rows);
        saveToVector(primaryKeys.getString("COLUMN_NAME"), columnNameVector, rows);
        final int key_seq = primaryKeys.getInt("KEY_SEQ");
        saveToVector(primaryKeys.wasNull() ? null : key_seq, keySequenceVector, rows);
        saveToVector(primaryKeys.getString("PK_NAME"), keyNameVector, rows);
      }

      try (final VectorSchemaRoot vectorSchemaRoot = new VectorSchemaRoot(vectors)) {
        vectorSchemaRoot.setRowCount(rows);

        listener.start(vectorSchemaRoot);
        listener.putNext();
      }
    } catch (SQLException e) {
      listener.error(e);
    } finally {
      listener.completed();
    }
  }

  @Override
  public FlightInfo getFlightInfoExportedKeys(
      final CommandGetExportedKeys request,
      final CallContext context,
      final FlightDescriptor descriptor) {
    return getFlightInfoForSchema(request, descriptor, Schemas.GET_EXPORTED_KEYS_SCHEMA);
  }

  @Override
  public void getStreamExportedKeys(
      final CommandGetExportedKeys command,
      final CallContext context,
      final ServerStreamListener listener) {
    String catalog = command.hasCatalog() ? command.getCatalog() : null;
    String schema = command.hasDbSchema() ? command.getDbSchema() : null;
    String table = command.getTable();

    try (Connection connection = DriverManager.getConnection(databaseUri);
        ResultSet keys = connection.getMetaData().getExportedKeys(catalog, schema, table);
        VectorSchemaRoot vectorSchemaRoot = createVectors(keys)) {
      listener.start(vectorSchemaRoot);
      listener.putNext();
    } catch (SQLException e) {
      listener.error(e);
    } finally {
      listener.completed();
    }
  }

  @Override
  public FlightInfo getFlightInfoImportedKeys(
      final CommandGetImportedKeys request,
      final CallContext context,
      final FlightDescriptor descriptor) {
    return getFlightInfoForSchema(request, descriptor, Schemas.GET_IMPORTED_KEYS_SCHEMA);
  }

  @Override
  public void getStreamImportedKeys(
      final CommandGetImportedKeys command,
      final CallContext context,
      final ServerStreamListener listener) {
    String catalog = command.hasCatalog() ? command.getCatalog() : null;
    String schema = command.hasDbSchema() ? command.getDbSchema() : null;
    String table = command.getTable();

    try (Connection connection = DriverManager.getConnection(databaseUri);
        ResultSet keys = connection.getMetaData().getImportedKeys(catalog, schema, table);
        VectorSchemaRoot vectorSchemaRoot = createVectors(keys)) {
      listener.start(vectorSchemaRoot);
      listener.putNext();
    } catch (final SQLException e) {
      listener.error(e);
    } finally {
      listener.completed();
    }
  }

  @Override
  public FlightInfo getFlightInfoCrossReference(
      CommandGetCrossReference request, CallContext context, FlightDescriptor descriptor) {
    return getFlightInfoForSchema(request, descriptor, Schemas.GET_CROSS_REFERENCE_SCHEMA);
  }

  @Override
  public void getStreamCrossReference(
      CommandGetCrossReference command, CallContext context, ServerStreamListener listener) {
    final String pkCatalog = command.hasPkCatalog() ? command.getPkCatalog() : null;
    final String pkSchema = command.hasPkDbSchema() ? command.getPkDbSchema() : null;
    final String fkCatalog = command.hasFkCatalog() ? command.getFkCatalog() : null;
    final String fkSchema = command.hasFkDbSchema() ? command.getFkDbSchema() : null;
    final String pkTable = command.getPkTable();
    final String fkTable = command.getFkTable();

    try (Connection connection = DriverManager.getConnection(databaseUri);
        ResultSet keys =
            connection
                .getMetaData()
                .getCrossReference(pkCatalog, pkSchema, pkTable, fkCatalog, fkSchema, fkTable);
        VectorSchemaRoot vectorSchemaRoot = createVectors(keys)) {
      listener.start(vectorSchemaRoot);
      listener.putNext();
    } catch (final SQLException e) {
      listener.error(e);
    } finally {
      listener.completed();
    }
  }

  private VectorSchemaRoot createVectors(ResultSet keys) throws SQLException {
    final VarCharVector pkCatalogNameVector = new VarCharVector("pk_catalog_name", rootAllocator);
    final VarCharVector pkSchemaNameVector = new VarCharVector("pk_db_schema_name", rootAllocator);
    final VarCharVector pkTableNameVector = new VarCharVector("pk_table_name", rootAllocator);
    final VarCharVector pkColumnNameVector = new VarCharVector("pk_column_name", rootAllocator);
    final VarCharVector fkCatalogNameVector = new VarCharVector("fk_catalog_name", rootAllocator);
    final VarCharVector fkSchemaNameVector = new VarCharVector("fk_db_schema_name", rootAllocator);
    final VarCharVector fkTableNameVector = new VarCharVector("fk_table_name", rootAllocator);
    final VarCharVector fkColumnNameVector = new VarCharVector("fk_column_name", rootAllocator);
    final IntVector keySequenceVector = new IntVector("key_sequence", rootAllocator);
    final VarCharVector fkKeyNameVector = new VarCharVector("fk_key_name", rootAllocator);
    final VarCharVector pkKeyNameVector = new VarCharVector("pk_key_name", rootAllocator);
    final UInt1Vector updateRuleVector = new UInt1Vector("update_rule", rootAllocator);
    final UInt1Vector deleteRuleVector = new UInt1Vector("delete_rule", rootAllocator);

    Map<FieldVector, String> vectorToColumnName = new HashMap<>();
    vectorToColumnName.put(pkCatalogNameVector, "PKTABLE_CAT");
    vectorToColumnName.put(pkSchemaNameVector, "PKTABLE_SCHEM");
    vectorToColumnName.put(pkTableNameVector, "PKTABLE_NAME");
    vectorToColumnName.put(pkColumnNameVector, "PKCOLUMN_NAME");
    vectorToColumnName.put(fkCatalogNameVector, "FKTABLE_CAT");
    vectorToColumnName.put(fkSchemaNameVector, "FKTABLE_SCHEM");
    vectorToColumnName.put(fkTableNameVector, "FKTABLE_NAME");
    vectorToColumnName.put(fkColumnNameVector, "FKCOLUMN_NAME");
    vectorToColumnName.put(keySequenceVector, "KEY_SEQ");
    vectorToColumnName.put(updateRuleVector, "UPDATE_RULE");
    vectorToColumnName.put(deleteRuleVector, "DELETE_RULE");
    vectorToColumnName.put(fkKeyNameVector, "FK_NAME");
    vectorToColumnName.put(pkKeyNameVector, "PK_NAME");

    final VectorSchemaRoot vectorSchemaRoot =
        VectorSchemaRoot.of(
            pkCatalogNameVector,
            pkSchemaNameVector,
            pkTableNameVector,
            pkColumnNameVector,
            fkCatalogNameVector,
            fkSchemaNameVector,
            fkTableNameVector,
            fkColumnNameVector,
            keySequenceVector,
            fkKeyNameVector,
            pkKeyNameVector,
            updateRuleVector,
            deleteRuleVector);

    vectorSchemaRoot.allocateNew();
    final int rowCount = saveToVectors(vectorToColumnName, keys, true);

    vectorSchemaRoot.setRowCount(rowCount);

    return vectorSchemaRoot;
  }

  @Override
  public void getStreamStatement(
      final TicketStatementQuery ticketStatementQuery,
      final CallContext context,
      final ServerStreamListener listener) {
    final ByteString handle = ticketStatementQuery.getStatementHandle();
    final StatementContext<Statement> statementContext =
        Objects.requireNonNull(statementLoadingCache.getIfPresent(handle));
    try (final ResultSet resultSet = statementContext.getStatement().getResultSet()) {
      final Schema schema = jdbcToArrowSchema(resultSet.getMetaData(), DEFAULT_CALENDAR);
      try (VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.create(schema, rootAllocator)) {
        final VectorLoader loader = new VectorLoader(vectorSchemaRoot);
        listener.start(vectorSchemaRoot);

        final ArrowVectorIterator iterator = sqlToArrowVectorIterator(resultSet, rootAllocator);
        while (iterator.hasNext()) {
          final VectorUnloader unloader = new VectorUnloader(iterator.next());
          loader.load(unloader.getRecordBatch());
          listener.putNext();
          vectorSchemaRoot.clear();
        }

        listener.putNext();
      }
    } catch (SQLException | IOException e) {
      LOGGER.error(format("Failed to getStreamPreparedStatement: <%s>.", e.getMessage()), e);
      listener.error(e);
    } finally {
      listener.completed();
      statementLoadingCache.invalidate(handle);
    }
  }

  protected <T extends Message> FlightInfo getFlightInfoForSchema(
      final T request, final FlightDescriptor descriptor, final Schema schema) {
    final Ticket ticket = new Ticket(pack(request).toByteArray());
    // TODO Support multiple endpoints.
    final List<FlightEndpoint> endpoints = singletonList(new FlightEndpoint(ticket, location));

    return new FlightInfo(schema, descriptor, endpoints, -1, -1);
  }

  private static class StatementRemovalListener<T extends Statement>
      implements RemovalListener<ByteString, StatementContext<T>> {
    @Override
    public void onRemoval(final RemovalNotification<ByteString, StatementContext<T>> notification) {
      try {
        AutoCloseables.close(notification.getValue());
      } catch (final Exception e) {
        // swallow
      }
    }
  }
}
