/*
 * 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;

import static java.util.Objects.isNull;
import static org.apache.arrow.flight.sql.impl.FlightSql.ActionBeginSavepointRequest;
import static org.apache.arrow.flight.sql.impl.FlightSql.ActionBeginSavepointResult;
import static org.apache.arrow.flight.sql.impl.FlightSql.ActionBeginTransactionRequest;
import static org.apache.arrow.flight.sql.impl.FlightSql.ActionBeginTransactionResult;
import static org.apache.arrow.flight.sql.impl.FlightSql.ActionCancelQueryRequest;
import static org.apache.arrow.flight.sql.impl.FlightSql.ActionCancelQueryResult;
import static org.apache.arrow.flight.sql.impl.FlightSql.ActionClosePreparedStatementRequest;
import static org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementRequest;
import static org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedSubstraitPlanRequest;
import static org.apache.arrow.flight.sql.impl.FlightSql.ActionEndSavepointRequest;
import static org.apache.arrow.flight.sql.impl.FlightSql.ActionEndTransactionRequest;
import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCatalogs;
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.CommandGetPrimaryKeys;
import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetSqlInfo;
import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTableTypes;
import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTables;
import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetXdbcTypeInfo;
import static org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementUpdate;
import static org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementQuery;
import static org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementSubstraitPlan;
import static org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementUpdate;
import static org.apache.arrow.flight.sql.impl.FlightSql.DoPutUpdateResult;
import static org.apache.arrow.flight.sql.impl.FlightSql.SqlInfo;

import com.google.protobuf.Any;
import com.google.protobuf.ByteString;
import com.google.protobuf.InvalidProtocolBufferException;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.nio.channels.Channels;
import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ExecutionException;
import java.util.function.Consumer;
import java.util.stream.Collectors;
import org.apache.arrow.flight.Action;
import org.apache.arrow.flight.CallOption;
import org.apache.arrow.flight.CallStatus;
import org.apache.arrow.flight.CancelFlightInfoRequest;
import org.apache.arrow.flight.CancelFlightInfoResult;
import org.apache.arrow.flight.CloseSessionRequest;
import org.apache.arrow.flight.CloseSessionResult;
import org.apache.arrow.flight.FlightClient;
import org.apache.arrow.flight.FlightDescriptor;
import org.apache.arrow.flight.FlightEndpoint;
import org.apache.arrow.flight.FlightInfo;
import org.apache.arrow.flight.FlightStream;
import org.apache.arrow.flight.GetSessionOptionsRequest;
import org.apache.arrow.flight.GetSessionOptionsResult;
import org.apache.arrow.flight.PutResult;
import org.apache.arrow.flight.RenewFlightEndpointRequest;
import org.apache.arrow.flight.Result;
import org.apache.arrow.flight.SchemaResult;
import org.apache.arrow.flight.SetSessionOptionsRequest;
import org.apache.arrow.flight.SetSessionOptionsResult;
import org.apache.arrow.flight.SyncPutListener;
import org.apache.arrow.flight.Ticket;
import org.apache.arrow.flight.sql.impl.FlightSql;
import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementResult;
import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementQuery;
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.util.TableRef;
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.util.Preconditions;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.ipc.ArrowReader;
import org.apache.arrow.vector.ipc.ReadChannel;
import org.apache.arrow.vector.ipc.message.MessageSerializer;
import org.apache.arrow.vector.types.pojo.Schema;

/** Flight client with Flight SQL semantics. */
public class FlightSqlClient implements AutoCloseable {
  private final FlightClient client;

  public FlightSqlClient(final FlightClient client) {
    this.client = Objects.requireNonNull(client, "Client cannot be null!");
  }

  /**
   * Execute a query on the server.
   *
   * @param query The query to execute.
   * @param options RPC-layer hints for this call.
   * @return a FlightInfo object representing the stream(s) to fetch.
   */
  public FlightInfo execute(final String query, final CallOption... options) {
    return execute(query, /*transaction*/ null, options);
  }

  /**
   * Execute a query on the server.
   *
   * @param query The query to execute.
   * @param transaction The transaction that this query is part of.
   * @param options RPC-layer hints for this call.
   * @return a FlightInfo object representing the stream(s) to fetch.
   */
  public FlightInfo execute(
      final String query, Transaction transaction, final CallOption... options) {
    final CommandStatementQuery.Builder builder =
        CommandStatementQuery.newBuilder().setQuery(query);
    if (transaction != null) {
      builder.setTransactionId(ByteString.copyFrom(transaction.getTransactionId()));
    }
    final FlightDescriptor descriptor =
        FlightDescriptor.command(Any.pack(builder.build()).toByteArray());
    return client.getInfo(descriptor, options);
  }

  /**
   * Execute a Substrait plan on the server.
   *
   * @param plan The Substrait plan to execute.
   * @param options RPC-layer hints for this call.
   * @return a FlightInfo object representing the stream(s) to fetch.
   */
  public FlightInfo executeSubstrait(SubstraitPlan plan, CallOption... options) {
    return executeSubstrait(plan, /*transaction*/ null, options);
  }

  /**
   * Execute a Substrait plan on the server.
   *
   * @param plan The Substrait plan to execute.
   * @param transaction The transaction that this query is part of.
   * @param options RPC-layer hints for this call.
   * @return a FlightInfo object representing the stream(s) to fetch.
   */
  public FlightInfo executeSubstrait(
      SubstraitPlan plan, Transaction transaction, CallOption... options) {
    final CommandStatementSubstraitPlan.Builder builder =
        CommandStatementSubstraitPlan.newBuilder();
    builder
        .getPlanBuilder()
        .setPlan(ByteString.copyFrom(plan.getPlan()))
        .setVersion(plan.getVersion());
    if (transaction != null) {
      builder.setTransactionId(ByteString.copyFrom(transaction.getTransactionId()));
    }
    final FlightDescriptor descriptor =
        FlightDescriptor.command(Any.pack(builder.build()).toByteArray());
    return client.getInfo(descriptor, options);
  }

  /** Get the schema of the result set of a query. */
  public SchemaResult getExecuteSchema(
      String query, Transaction transaction, CallOption... options) {
    final CommandStatementQuery.Builder builder = CommandStatementQuery.newBuilder();
    builder.setQuery(query);
    if (transaction != null) {
      builder.setTransactionId(ByteString.copyFrom(transaction.getTransactionId()));
    }
    final FlightDescriptor descriptor =
        FlightDescriptor.command(Any.pack(builder.build()).toByteArray());
    return client.getSchema(descriptor, options);
  }

  /** Get the schema of the result set of a query. */
  public SchemaResult getExecuteSchema(String query, CallOption... options) {
    return getExecuteSchema(query, /*transaction*/ null, options);
  }

  /** Get the schema of the result set of a Substrait plan. */
  public SchemaResult getExecuteSubstraitSchema(
      SubstraitPlan plan, Transaction transaction, final CallOption... options) {
    final CommandStatementSubstraitPlan.Builder builder =
        CommandStatementSubstraitPlan.newBuilder();
    builder
        .getPlanBuilder()
        .setPlan(ByteString.copyFrom(plan.getPlan()))
        .setVersion(plan.getVersion());
    if (transaction != null) {
      builder.setTransactionId(ByteString.copyFrom(transaction.getTransactionId()));
    }
    final FlightDescriptor descriptor =
        FlightDescriptor.command(Any.pack(builder.build()).toByteArray());
    return client.getSchema(descriptor, options);
  }

  /** Get the schema of the result set of a Substrait plan. */
  public SchemaResult getExecuteSubstraitSchema(
      SubstraitPlan substraitPlan, final CallOption... options) {
    return getExecuteSubstraitSchema(substraitPlan, /*transaction*/ null, options);
  }

  /**
   * Execute a bulk ingest on the server.
   *
   * @param data data to be ingested
   * @param ingestOptions options for the ingest request.
   * @param options RPC-layer hints for this call.
   * @return the number of rows affected.
   */
  public long executeIngest(
      final VectorSchemaRoot data,
      final ExecuteIngestOptions ingestOptions,
      final CallOption... options) {
    return executeIngest(data, ingestOptions, /*transaction*/ null, options);
  }

  /**
   * Execute a bulk ingest on the server.
   *
   * @param dataReader data stream to be ingested
   * @param ingestOptions options for the ingest request.
   * @param options RPC-layer hints for this call.
   * @return the number of rows affected.
   */
  public long executeIngest(
      final ArrowReader dataReader,
      final ExecuteIngestOptions ingestOptions,
      final CallOption... options) {
    return executeIngest(dataReader, ingestOptions, /*transaction*/ null, options);
  }

  /**
   * Execute a bulk ingest on the server.
   *
   * @param data data to be ingested
   * @param ingestOptions options for the ingest request.
   * @param transaction The transaction that this ingest request is part of.
   * @param options RPC-layer hints for this call.
   * @return the number of rows affected.
   */
  public long executeIngest(
      final VectorSchemaRoot data,
      final ExecuteIngestOptions ingestOptions,
      Transaction transaction,
      final CallOption... options) {
    return executeIngest(
        data, ingestOptions, transaction, FlightClient.ClientStreamListener::putNext, options);
  }

  /**
   * Execute a bulk ingest on the server.
   *
   * @param dataReader data stream to be ingested
   * @param ingestOptions options for the ingest request.
   * @param transaction The transaction that this ingest request is part of.
   * @param options RPC-layer hints for this call.
   * @return the number of rows affected.
   */
  public long executeIngest(
      final ArrowReader dataReader,
      final ExecuteIngestOptions ingestOptions,
      Transaction transaction,
      final CallOption... options) {

    try {
      return executeIngest(
          dataReader.getVectorSchemaRoot(),
          ingestOptions,
          transaction,
          listener -> {
            while (true) {
              try {
                if (!dataReader.loadNextBatch()) {
                  break;
                }
              } catch (IOException e) {
                throw CallStatus.UNKNOWN.withCause(e).toRuntimeException();
              }
              listener.putNext();
            }
          },
          options);
    } catch (IOException e) {
      throw CallStatus.UNKNOWN.withCause(e).toRuntimeException();
    }
  }

  private long executeIngest(
      final VectorSchemaRoot data,
      final ExecuteIngestOptions ingestOptions,
      final Transaction transaction,
      final Consumer<FlightClient.ClientStreamListener> dataPutter,
      final CallOption... options) {
    try {
      final CommandStatementIngest.Builder builder = CommandStatementIngest.newBuilder();
      if (transaction != null) {
        builder.setTransactionId(ByteString.copyFrom(transaction.getTransactionId()));
      }
      ingestOptions.updateCommandBuilder(builder);

      final FlightDescriptor descriptor =
          FlightDescriptor.command(Any.pack(builder.build()).toByteArray());
      try (final SyncPutListener putListener = new SyncPutListener()) {

        final FlightClient.ClientStreamListener listener =
            client.startPut(descriptor, data, putListener, options);
        dataPutter.accept(listener);
        listener.completed();
        listener.getResult();

        try (final PutResult result = putListener.read()) {
          final DoPutUpdateResult doPutUpdateResult =
              DoPutUpdateResult.parseFrom(result.getApplicationMetadata().nioBuffer());
          return doPutUpdateResult.getRecordCount();
        }
      }
    } catch (final InterruptedException e) {
      throw CallStatus.CANCELLED.withCause(e).toRuntimeException();
    } catch (final ExecutionException e) {
      throw CallStatus.CANCELLED.withCause(e.getCause()).toRuntimeException();
    } catch (final InvalidProtocolBufferException e) {
      throw CallStatus.INTERNAL.withCause(e).toRuntimeException();
    }
  }

  /**
   * Execute an update query on the server.
   *
   * @param query The query to execute.
   * @param options RPC-layer hints for this call.
   * @return the number of rows affected.
   */
  public long executeUpdate(final String query, final CallOption... options) {
    return executeUpdate(query, /*transaction*/ null, options);
  }

  /**
   * Execute an update query on the server.
   *
   * @param query The query to execute.
   * @param transaction The transaction that this query is part of.
   * @param options RPC-layer hints for this call.
   * @return the number of rows affected.
   */
  public long executeUpdate(
      final String query, Transaction transaction, final CallOption... options) {
    final CommandStatementUpdate.Builder builder =
        CommandStatementUpdate.newBuilder().setQuery(query);
    if (transaction != null) {
      builder.setTransactionId(ByteString.copyFrom(transaction.getTransactionId()));
    }

    final FlightDescriptor descriptor =
        FlightDescriptor.command(Any.pack(builder.build()).toByteArray());
    try (final SyncPutListener putListener = new SyncPutListener()) {
      final FlightClient.ClientStreamListener listener =
          client.startPut(descriptor, VectorSchemaRoot.of(), putListener, options);
      try (final PutResult result = putListener.read()) {
        final DoPutUpdateResult doPutUpdateResult =
            DoPutUpdateResult.parseFrom(result.getApplicationMetadata().nioBuffer());
        return doPutUpdateResult.getRecordCount();
      } finally {
        listener.getResult();
      }
    } catch (final InterruptedException e) {
      throw CallStatus.CANCELLED.withCause(e).toRuntimeException();
    } catch (final ExecutionException e) {
      throw CallStatus.CANCELLED.withCause(e.getCause()).toRuntimeException();
    } catch (final InvalidProtocolBufferException e) {
      throw CallStatus.INTERNAL.withCause(e).toRuntimeException();
    }
  }

  /**
   * Execute an update query on the server.
   *
   * @param plan The Substrait plan to execute.
   * @param options RPC-layer hints for this call.
   * @return the number of rows affected.
   */
  public long executeSubstraitUpdate(SubstraitPlan plan, CallOption... options) {
    return executeSubstraitUpdate(plan, /*transaction*/ null, options);
  }

  /**
   * Execute an update query on the server.
   *
   * @param plan The Substrait plan to execute.
   * @param transaction The transaction that this query is part of.
   * @param options RPC-layer hints for this call.
   * @return the number of rows affected.
   */
  public long executeSubstraitUpdate(
      SubstraitPlan plan, Transaction transaction, CallOption... options) {
    final CommandStatementSubstraitPlan.Builder builder =
        CommandStatementSubstraitPlan.newBuilder();
    builder
        .getPlanBuilder()
        .setPlan(ByteString.copyFrom(plan.getPlan()))
        .setVersion(plan.getVersion());
    if (transaction != null) {
      builder.setTransactionId(ByteString.copyFrom(transaction.getTransactionId()));
    }

    final FlightDescriptor descriptor =
        FlightDescriptor.command(Any.pack(builder.build()).toByteArray());
    try (final SyncPutListener putListener = new SyncPutListener()) {
      final FlightClient.ClientStreamListener listener =
          client.startPut(descriptor, VectorSchemaRoot.of(), putListener, options);
      try (final PutResult result = putListener.read()) {
        final DoPutUpdateResult doPutUpdateResult =
            DoPutUpdateResult.parseFrom(result.getApplicationMetadata().nioBuffer());
        return doPutUpdateResult.getRecordCount();
      } finally {
        listener.getResult();
      }
    } catch (final InterruptedException e) {
      throw CallStatus.CANCELLED.withCause(e).toRuntimeException();
    } catch (final ExecutionException e) {
      throw CallStatus.CANCELLED.withCause(e.getCause()).toRuntimeException();
    } catch (final InvalidProtocolBufferException e) {
      throw CallStatus.INTERNAL.withCause(e).toRuntimeException();
    }
  }

  /**
   * Request a list of catalogs.
   *
   * @param options RPC-layer hints for this call.
   * @return a FlightInfo object representing the stream(s) to fetch.
   */
  public FlightInfo getCatalogs(final CallOption... options) {
    final CommandGetCatalogs.Builder builder = CommandGetCatalogs.newBuilder();
    final FlightDescriptor descriptor =
        FlightDescriptor.command(Any.pack(builder.build()).toByteArray());
    return client.getInfo(descriptor, options);
  }

  /**
   * Get the schema of {@link #getCatalogs(CallOption...)} from the server.
   *
   * <p>Should be identical to {@link FlightSqlProducer.Schemas#GET_CATALOGS_SCHEMA}.
   */
  public SchemaResult getCatalogsSchema(final CallOption... options) {
    final CommandGetCatalogs command = CommandGetCatalogs.getDefaultInstance();
    final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray());
    return client.getSchema(descriptor, options);
  }

  /**
   * Request a list of schemas.
   *
   * @param catalog The catalog.
   * @param dbSchemaFilterPattern The schema filter pattern.
   * @param options RPC-layer hints for this call.
   * @return a FlightInfo object representing the stream(s) to fetch.
   */
  public FlightInfo getSchemas(
      final String catalog, final String dbSchemaFilterPattern, final CallOption... options) {
    final CommandGetDbSchemas.Builder builder = CommandGetDbSchemas.newBuilder();

    if (catalog != null) {
      builder.setCatalog(catalog);
    }

    if (dbSchemaFilterPattern != null) {
      builder.setDbSchemaFilterPattern(dbSchemaFilterPattern);
    }

    final FlightDescriptor descriptor =
        FlightDescriptor.command(Any.pack(builder.build()).toByteArray());
    return client.getInfo(descriptor, options);
  }

  /**
   * Get the schema of {@link #getSchemas(String, String, CallOption...)} from the server.
   *
   * <p>Should be identical to {@link FlightSqlProducer.Schemas#GET_SCHEMAS_SCHEMA}.
   */
  public SchemaResult getSchemasSchema(final CallOption... options) {
    final CommandGetDbSchemas command = CommandGetDbSchemas.getDefaultInstance();
    final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray());
    return client.getSchema(descriptor, options);
  }

  /**
   * Get schema for a stream.
   *
   * @param descriptor The descriptor for the stream.
   * @param options RPC-layer hints for this call.
   */
  public SchemaResult getSchema(FlightDescriptor descriptor, CallOption... options) {
    return client.getSchema(descriptor, options);
  }

  /**
   * Retrieve a stream from the server.
   *
   * @param ticket The ticket granting access to the data stream.
   * @param options RPC-layer hints for this call.
   */
  public FlightStream getStream(Ticket ticket, CallOption... options) {
    return client.getStream(ticket, options);
  }

  /**
   * Request a set of Flight SQL metadata.
   *
   * @param info The set of metadata to retrieve. None to retrieve all metadata.
   * @return a FlightInfo object representing the stream(s) to fetch.
   */
  public FlightInfo getSqlInfo(final SqlInfo... info) {
    return getSqlInfo(info, new CallOption[0]);
  }

  /**
   * Request a set of Flight SQL metadata.
   *
   * @param info The set of metadata to retrieve. None to retrieve all metadata.
   * @param options RPC-layer hints for this call.
   * @return a FlightInfo object representing the stream(s) to fetch.
   */
  public FlightInfo getSqlInfo(final SqlInfo[] info, final CallOption... options) {
    final int[] infoNumbers = Arrays.stream(info).mapToInt(SqlInfo::getNumber).toArray();
    return getSqlInfo(infoNumbers, options);
  }

  /**
   * Request a set of Flight SQL metadata. Use this method if you would like to retrieve custom
   * metadata, where the custom metadata key values start from 10_000.
   *
   * @param info The set of metadata to retrieve. None to retrieve all metadata.
   * @param options RPC-layer hints for this call.
   * @return a FlightInfo object representing the stream(s) to fetch.
   */
  public FlightInfo getSqlInfo(final int[] info, final CallOption... options) {
    return getSqlInfo(Arrays.stream(info).boxed().collect(Collectors.toList()), options);
  }

  /**
   * Request a set of Flight SQL metadata. Use this method if you would like to retrieve custom
   * metadata, where the custom metadata key values start from 10_000.
   *
   * @param info The set of metadata to retrieve. None to retrieve all metadata.
   * @param options RPC-layer hints for this call.
   * @return a FlightInfo object representing the stream(s) to fetch.
   */
  public FlightInfo getSqlInfo(final Iterable<Integer> info, final CallOption... options) {
    final CommandGetSqlInfo.Builder builder = CommandGetSqlInfo.newBuilder();
    builder.addAllInfo(info);
    final FlightDescriptor descriptor =
        FlightDescriptor.command(Any.pack(builder.build()).toByteArray());
    return client.getInfo(descriptor, options);
  }

  /**
   * Get the schema of {@link #getSqlInfo(SqlInfo...)} from the server.
   *
   * <p>Should be identical to {@link FlightSqlProducer.Schemas#GET_SQL_INFO_SCHEMA}.
   */
  public SchemaResult getSqlInfoSchema(final CallOption... options) {
    final CommandGetSqlInfo command = CommandGetSqlInfo.getDefaultInstance();
    final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray());
    return client.getSchema(descriptor, options);
  }

  /**
   * Request the information about the data types supported related to a filter data type.
   *
   * @param dataType the data type to be used as filter.
   * @param options RPC-layer hints for this call.
   * @return a FlightInfo object representing the stream(s) to fetch.
   */
  public FlightInfo getXdbcTypeInfo(final int dataType, final CallOption... options) {
    final CommandGetXdbcTypeInfo.Builder builder = CommandGetXdbcTypeInfo.newBuilder();

    builder.setDataType(dataType);

    final FlightDescriptor descriptor =
        FlightDescriptor.command(Any.pack(builder.build()).toByteArray());
    return client.getInfo(descriptor, options);
  }

  /**
   * Request the information about all the data types supported.
   *
   * @param options RPC-layer hints for this call.
   * @return a FlightInfo object representing the stream(s) to fetch.
   */
  public FlightInfo getXdbcTypeInfo(final CallOption... options) {
    final CommandGetXdbcTypeInfo.Builder builder = CommandGetXdbcTypeInfo.newBuilder();

    final FlightDescriptor descriptor =
        FlightDescriptor.command(Any.pack(builder.build()).toByteArray());
    return client.getInfo(descriptor, options);
  }

  /**
   * Get the schema of {@link #getXdbcTypeInfo(CallOption...)} from the server.
   *
   * <p>Should be identical to {@link FlightSqlProducer.Schemas#GET_TYPE_INFO_SCHEMA}.
   */
  public SchemaResult getXdbcTypeInfoSchema(final CallOption... options) {
    final CommandGetXdbcTypeInfo command = CommandGetXdbcTypeInfo.getDefaultInstance();
    final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray());
    return client.getSchema(descriptor, options);
  }

  /**
   * Request a list of tables.
   *
   * @param catalog The catalog.
   * @param dbSchemaFilterPattern The schema filter pattern.
   * @param tableFilterPattern The table filter pattern.
   * @param tableTypes The table types to include.
   * @param includeSchema True to include the schema upon return, false to not include the schema.
   * @param options RPC-layer hints for this call.
   * @return a FlightInfo object representing the stream(s) to fetch.
   */
  public FlightInfo getTables(
      final String catalog,
      final String dbSchemaFilterPattern,
      final String tableFilterPattern,
      final List<String> tableTypes,
      final boolean includeSchema,
      final CallOption... options) {
    final CommandGetTables.Builder builder = CommandGetTables.newBuilder();

    if (catalog != null) {
      builder.setCatalog(catalog);
    }

    if (dbSchemaFilterPattern != null) {
      builder.setDbSchemaFilterPattern(dbSchemaFilterPattern);
    }

    if (tableFilterPattern != null) {
      builder.setTableNameFilterPattern(tableFilterPattern);
    }

    if (tableTypes != null) {
      builder.addAllTableTypes(tableTypes);
    }
    builder.setIncludeSchema(includeSchema);

    final FlightDescriptor descriptor =
        FlightDescriptor.command(Any.pack(builder.build()).toByteArray());
    return client.getInfo(descriptor, options);
  }

  /**
   * Get the schema of {@link #getTables(String, String, String, List, boolean, CallOption...)} from
   * the server.
   *
   * <p>Should be identical to {@link FlightSqlProducer.Schemas#GET_TABLES_SCHEMA} or {@link
   * FlightSqlProducer.Schemas#GET_TABLES_SCHEMA_NO_SCHEMA}.
   */
  public SchemaResult getTablesSchema(boolean includeSchema, final CallOption... options) {
    final CommandGetTables command =
        CommandGetTables.newBuilder().setIncludeSchema(includeSchema).build();
    final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray());
    return client.getSchema(descriptor, options);
  }

  /**
   * Request the primary keys for a table.
   *
   * @param tableRef An object which hold info about catalog, dbSchema and table.
   * @param options RPC-layer hints for this call.
   * @return a FlightInfo object representing the stream(s) to fetch.
   */
  public FlightInfo getPrimaryKeys(final TableRef tableRef, final CallOption... options) {
    final CommandGetPrimaryKeys.Builder builder = CommandGetPrimaryKeys.newBuilder();

    if (tableRef.getCatalog() != null) {
      builder.setCatalog(tableRef.getCatalog());
    }

    if (tableRef.getDbSchema() != null) {
      builder.setDbSchema(tableRef.getDbSchema());
    }

    Objects.requireNonNull(tableRef.getTable());
    builder.setTable(tableRef.getTable());

    final FlightDescriptor descriptor =
        FlightDescriptor.command(Any.pack(builder.build()).toByteArray());
    return client.getInfo(descriptor, options);
  }

  /**
   * Get the schema of {@link #getPrimaryKeys(TableRef, CallOption...)} from the server.
   *
   * <p>Should be identical to {@link FlightSqlProducer.Schemas#GET_PRIMARY_KEYS_SCHEMA}.
   */
  public SchemaResult getPrimaryKeysSchema(final CallOption... options) {
    final CommandGetPrimaryKeys command = CommandGetPrimaryKeys.getDefaultInstance();
    final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray());
    return client.getSchema(descriptor, options);
  }

  /**
   * Retrieves a description about the foreign key columns that reference the primary key columns of
   * the given table.
   *
   * @param tableRef An object which hold info about catalog, dbSchema and table.
   * @param options RPC-layer hints for this call.
   * @return a FlightInfo object representing the stream(s) to fetch.
   */
  public FlightInfo getExportedKeys(final TableRef tableRef, final CallOption... options) {
    Objects.requireNonNull(tableRef.getTable(), "Table cannot be null.");

    final CommandGetExportedKeys.Builder builder = CommandGetExportedKeys.newBuilder();

    if (tableRef.getCatalog() != null) {
      builder.setCatalog(tableRef.getCatalog());
    }

    if (tableRef.getDbSchema() != null) {
      builder.setDbSchema(tableRef.getDbSchema());
    }

    Objects.requireNonNull(tableRef.getTable());
    builder.setTable(tableRef.getTable());

    final FlightDescriptor descriptor =
        FlightDescriptor.command(Any.pack(builder.build()).toByteArray());
    return client.getInfo(descriptor, options);
  }

  /**
   * Get the schema of {@link #getExportedKeys(TableRef, CallOption...)} from the server.
   *
   * <p>Should be identical to {@link FlightSqlProducer.Schemas#GET_EXPORTED_KEYS_SCHEMA}.
   */
  public SchemaResult getExportedKeysSchema(final CallOption... options) {
    final CommandGetExportedKeys command = CommandGetExportedKeys.getDefaultInstance();
    final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray());
    return client.getSchema(descriptor, options);
  }

  /**
   * Retrieves the foreign key columns for the given table.
   *
   * @param tableRef An object which hold info about catalog, dbSchema and table.
   * @param options RPC-layer hints for this call.
   * @return a FlightInfo object representing the stream(s) to fetch.
   */
  public FlightInfo getImportedKeys(final TableRef tableRef, final CallOption... options) {
    Objects.requireNonNull(tableRef.getTable(), "Table cannot be null.");

    final CommandGetImportedKeys.Builder builder = CommandGetImportedKeys.newBuilder();

    if (tableRef.getCatalog() != null) {
      builder.setCatalog(tableRef.getCatalog());
    }

    if (tableRef.getDbSchema() != null) {
      builder.setDbSchema(tableRef.getDbSchema());
    }

    Objects.requireNonNull(tableRef.getTable());
    builder.setTable(tableRef.getTable());

    final FlightDescriptor descriptor =
        FlightDescriptor.command(Any.pack(builder.build()).toByteArray());
    return client.getInfo(descriptor, options);
  }

  /**
   * Get the schema of {@link #getImportedKeys(TableRef, CallOption...)} from the server.
   *
   * <p>Should be identical to {@link FlightSqlProducer.Schemas#GET_IMPORTED_KEYS_SCHEMA}.
   */
  public SchemaResult getImportedKeysSchema(final CallOption... options) {
    final CommandGetImportedKeys command = CommandGetImportedKeys.getDefaultInstance();
    final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray());
    return client.getSchema(descriptor, options);
  }

  /**
   * Retrieves a description of the foreign key columns that reference the given table's primary key
   * columns (the foreign keys exported by a table).
   *
   * @param pkTableRef An object which hold info about catalog, dbSchema and table from a primary
   *     table.
   * @param fkTableRef An object which hold info about catalog, dbSchema and table from a foreign
   *     table.
   * @param options RPC-layer hints for this call.
   * @return a FlightInfo object representing the stream(s) to fetch.
   */
  public FlightInfo getCrossReference(
      final TableRef pkTableRef, final TableRef fkTableRef, final CallOption... options) {
    Objects.requireNonNull(pkTableRef.getTable(), "Parent Table cannot be null.");
    Objects.requireNonNull(fkTableRef.getTable(), "Foreign Table cannot be null.");

    final CommandGetCrossReference.Builder builder = CommandGetCrossReference.newBuilder();

    if (pkTableRef.getCatalog() != null) {
      builder.setPkCatalog(pkTableRef.getCatalog());
    }

    if (pkTableRef.getDbSchema() != null) {
      builder.setPkDbSchema(pkTableRef.getDbSchema());
    }

    if (fkTableRef.getCatalog() != null) {
      builder.setFkCatalog(fkTableRef.getCatalog());
    }

    if (fkTableRef.getDbSchema() != null) {
      builder.setFkDbSchema(fkTableRef.getDbSchema());
    }

    builder.setPkTable(pkTableRef.getTable());
    builder.setFkTable(fkTableRef.getTable());

    final FlightDescriptor descriptor =
        FlightDescriptor.command(Any.pack(builder.build()).toByteArray());
    return client.getInfo(descriptor, options);
  }

  /**
   * Get the schema of {@link #getCrossReference(TableRef, TableRef, CallOption...)} from the
   * server.
   *
   * <p>Should be identical to {@link FlightSqlProducer.Schemas#GET_CROSS_REFERENCE_SCHEMA}.
   */
  public SchemaResult getCrossReferenceSchema(final CallOption... options) {
    final CommandGetCrossReference command = CommandGetCrossReference.getDefaultInstance();
    final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray());
    return client.getSchema(descriptor, options);
  }

  /**
   * Request a list of table types.
   *
   * @param options RPC-layer hints for this call.
   * @return a FlightInfo object representing the stream(s) to fetch.
   */
  public FlightInfo getTableTypes(final CallOption... options) {
    final CommandGetTableTypes.Builder builder = CommandGetTableTypes.newBuilder();
    final FlightDescriptor descriptor =
        FlightDescriptor.command(Any.pack(builder.build()).toByteArray());
    return client.getInfo(descriptor, options);
  }

  /**
   * Get the schema of {@link #getTableTypes(CallOption...)} from the server.
   *
   * <p>Should be identical to {@link FlightSqlProducer.Schemas#GET_TABLE_TYPES_SCHEMA}.
   */
  public SchemaResult getTableTypesSchema(final CallOption... options) {
    final CommandGetTableTypes command = CommandGetTableTypes.getDefaultInstance();
    final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray());
    return client.getSchema(descriptor, options);
  }

  /**
   * Create a prepared statement for a SQL query on the server.
   *
   * @param query The query to prepare.
   * @param options RPC-layer hints for this call.
   * @return The representation of the prepared statement which exists on the server.
   */
  public PreparedStatement prepare(String query, CallOption... options) {
    return prepare(query, /*transaction*/ null, options);
  }

  /**
   * Create a prepared statement for a SQL query on the server.
   *
   * @param query The query to prepare.
   * @param transaction The transaction that this query is part of.
   * @param options RPC-layer hints for this call.
   * @return The representation of the prepared statement which exists on the server.
   */
  public PreparedStatement prepare(String query, Transaction transaction, CallOption... options) {
    ActionCreatePreparedStatementRequest.Builder builder =
        ActionCreatePreparedStatementRequest.newBuilder().setQuery(query);
    if (transaction != null) {
      builder.setTransactionId(ByteString.copyFrom(transaction.getTransactionId()));
    }
    return new PreparedStatement(
        client,
        new Action(
            FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_STATEMENT.getType(),
            Any.pack(builder.build()).toByteArray()),
        options);
  }

  /**
   * Create a prepared statement for a Substrait plan on the server.
   *
   * @param plan The query to prepare.
   * @param options RPC-layer hints for this call.
   * @return The representation of the prepared statement which exists on the server.
   */
  public PreparedStatement prepare(SubstraitPlan plan, CallOption... options) {
    return prepare(plan, /*transaction*/ null, options);
  }

  /**
   * Create a prepared statement for a Substrait plan on the server.
   *
   * @param plan The query to prepare.
   * @param transaction The transaction that this query is part of.
   * @param options RPC-layer hints for this call.
   * @return The representation of the prepared statement which exists on the server.
   */
  public PreparedStatement prepare(
      SubstraitPlan plan, Transaction transaction, CallOption... options) {
    ActionCreatePreparedSubstraitPlanRequest.Builder builder =
        ActionCreatePreparedSubstraitPlanRequest.newBuilder();
    builder
        .getPlanBuilder()
        .setPlan(ByteString.copyFrom(plan.getPlan()))
        .setVersion(plan.getVersion());
    if (transaction != null) {
      builder.setTransactionId(ByteString.copyFrom(transaction.getTransactionId()));
    }
    return new PreparedStatement(
        client,
        new Action(
            FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_SUBSTRAIT_PLAN.getType(),
            Any.pack(builder.build()).toByteArray()),
        options);
  }

  /** Begin a transaction. */
  public Transaction beginTransaction(CallOption... options) {
    final Action action =
        new Action(
            FlightSqlUtils.FLIGHT_SQL_BEGIN_TRANSACTION.getType(),
            Any.pack(ActionBeginTransactionRequest.getDefaultInstance()).toByteArray());
    final Iterator<Result> preparedStatementResults = client.doAction(action, options);
    final ActionBeginTransactionResult result =
        FlightSqlUtils.unpackAndParseOrThrow(
            preparedStatementResults.next().getBody(), ActionBeginTransactionResult.class);
    preparedStatementResults.forEachRemaining((ignored) -> {});
    if (result.getTransactionId().isEmpty()) {
      throw CallStatus.INTERNAL
          .withDescription("Server returned an empty transaction ID")
          .toRuntimeException();
    }
    return new Transaction(result.getTransactionId().toByteArray());
  }

  /** Create a savepoint within a transaction. */
  public Savepoint beginSavepoint(Transaction transaction, String name, CallOption... options) {
    Preconditions.checkArgument(
        transaction.getTransactionId().length != 0, "Transaction must be initialized");
    ActionBeginSavepointRequest request =
        ActionBeginSavepointRequest.newBuilder()
            .setTransactionId(ByteString.copyFrom(transaction.getTransactionId()))
            .setName(name)
            .build();
    final Action action =
        new Action(
            FlightSqlUtils.FLIGHT_SQL_BEGIN_SAVEPOINT.getType(), Any.pack(request).toByteArray());
    final Iterator<Result> preparedStatementResults = client.doAction(action, options);
    final ActionBeginSavepointResult result =
        FlightSqlUtils.unpackAndParseOrThrow(
            preparedStatementResults.next().getBody(), ActionBeginSavepointResult.class);
    preparedStatementResults.forEachRemaining((ignored) -> {});
    if (result.getSavepointId().isEmpty()) {
      throw CallStatus.INTERNAL
          .withDescription("Server returned an empty transaction ID")
          .toRuntimeException();
    }
    return new Savepoint(result.getSavepointId().toByteArray());
  }

  /** Commit a transaction. */
  public void commit(Transaction transaction, CallOption... options) {
    Preconditions.checkArgument(
        transaction.getTransactionId().length != 0, "Transaction must be initialized");
    ActionEndTransactionRequest request =
        ActionEndTransactionRequest.newBuilder()
            .setTransactionId(ByteString.copyFrom(transaction.getTransactionId()))
            .setActionValue(
                ActionEndTransactionRequest.EndTransaction.END_TRANSACTION_COMMIT.getNumber())
            .build();
    final Action action =
        new Action(
            FlightSqlUtils.FLIGHT_SQL_END_TRANSACTION.getType(), Any.pack(request).toByteArray());
    final Iterator<Result> preparedStatementResults = client.doAction(action, options);
    preparedStatementResults.forEachRemaining((ignored) -> {});
  }

  /** Release a savepoint. */
  public void release(Savepoint savepoint, CallOption... options) {
    Preconditions.checkArgument(
        savepoint.getSavepointId().length != 0, "Savepoint must be initialized");
    ActionEndSavepointRequest request =
        ActionEndSavepointRequest.newBuilder()
            .setSavepointId(ByteString.copyFrom(savepoint.getSavepointId()))
            .setActionValue(
                ActionEndSavepointRequest.EndSavepoint.END_SAVEPOINT_RELEASE.getNumber())
            .build();
    final Action action =
        new Action(
            FlightSqlUtils.FLIGHT_SQL_END_SAVEPOINT.getType(), Any.pack(request).toByteArray());
    final Iterator<Result> preparedStatementResults = client.doAction(action, options);
    preparedStatementResults.forEachRemaining((ignored) -> {});
  }

  /** Rollback a transaction. */
  public void rollback(Transaction transaction, CallOption... options) {
    Preconditions.checkArgument(
        transaction.getTransactionId().length != 0, "Transaction must be initialized");
    ActionEndTransactionRequest request =
        ActionEndTransactionRequest.newBuilder()
            .setTransactionId(ByteString.copyFrom(transaction.getTransactionId()))
            .setActionValue(
                ActionEndTransactionRequest.EndTransaction.END_TRANSACTION_ROLLBACK.getNumber())
            .build();
    final Action action =
        new Action(
            FlightSqlUtils.FLIGHT_SQL_END_TRANSACTION.getType(), Any.pack(request).toByteArray());
    final Iterator<Result> preparedStatementResults = client.doAction(action, options);
    preparedStatementResults.forEachRemaining((ignored) -> {});
  }

  /** Rollback to a savepoint. */
  public void rollback(Savepoint savepoint, CallOption... options) {
    Preconditions.checkArgument(
        savepoint.getSavepointId().length != 0, "Savepoint must be initialized");
    ActionEndSavepointRequest request =
        ActionEndSavepointRequest.newBuilder()
            .setSavepointId(ByteString.copyFrom(savepoint.getSavepointId()))
            .setActionValue(
                ActionEndSavepointRequest.EndSavepoint.END_SAVEPOINT_RELEASE.getNumber())
            .build();
    final Action action =
        new Action(
            FlightSqlUtils.FLIGHT_SQL_END_SAVEPOINT.getType(), Any.pack(request).toByteArray());
    final Iterator<Result> preparedStatementResults = client.doAction(action, options);
    preparedStatementResults.forEachRemaining((ignored) -> {});
  }

  /**
   * Cancel execution of a distributed query.
   *
   * @param request The query to cancel.
   * @param options Call options.
   * @return The server response.
   */
  public CancelFlightInfoResult cancelFlightInfo(
      CancelFlightInfoRequest request, CallOption... options) {
    return client.cancelFlightInfo(request, options);
  }

  /**
   * Explicitly cancel a running query.
   *
   * <p>This lets a single client explicitly cancel work, no matter how many clients are
   * involved/whether the query is distributed or not, given server support. The
   * transaction/statement is not rolled back; it is the application's job to commit or rollback as
   * appropriate. This only indicates the client no longer wishes to read the remainder of the query
   * results or continue submitting data.
   *
   * @deprecated Prefer {@link #cancelFlightInfo}.
   */
  @Deprecated
  public CancelResult cancelQuery(FlightInfo info, CallOption... options) {
    ActionCancelQueryRequest request =
        ActionCancelQueryRequest.newBuilder()
            .setInfo(ByteString.copyFrom(info.serialize()))
            .build();
    final Action action =
        new Action(
            FlightSqlUtils.FLIGHT_SQL_CANCEL_QUERY.getType(), Any.pack(request).toByteArray());
    final Iterator<Result> preparedStatementResults = client.doAction(action, options);
    final ActionCancelQueryResult result =
        FlightSqlUtils.unpackAndParseOrThrow(
            preparedStatementResults.next().getBody(), ActionCancelQueryResult.class);
    preparedStatementResults.forEachRemaining((ignored) -> {});
    switch (result.getResult()) {
      case CANCEL_RESULT_UNSPECIFIED:
        return CancelResult.UNSPECIFIED;
      case CANCEL_RESULT_CANCELLED:
        return CancelResult.CANCELLED;
      case CANCEL_RESULT_CANCELLING:
        return CancelResult.CANCELLING;
      case CANCEL_RESULT_NOT_CANCELLABLE:
        return CancelResult.NOT_CANCELLABLE;
      case UNRECOGNIZED:
      default:
        throw CallStatus.INTERNAL
            .withDescription("Unknown result: " + result.getResult())
            .toRuntimeException();
    }
  }

  /**
   * Request the server to extend the lifetime of a query result set.
   *
   * @param request The result set partition.
   * @param options Call options.
   * @return The new endpoint with an updated expiration time.
   */
  public FlightEndpoint renewFlightEndpoint(
      RenewFlightEndpointRequest request, CallOption... options) {
    return client.renewFlightEndpoint(request, options);
  }

  public SetSessionOptionsResult setSessionOptions(
      SetSessionOptionsRequest request, CallOption... options) {
    return client.setSessionOptions(request, options);
  }

  public GetSessionOptionsResult getSessionOptions(
      GetSessionOptionsRequest request, CallOption... options) {
    return client.getSessionOptions(request, options);
  }

  public CloseSessionResult closeSession(CloseSessionRequest request, CallOption... options) {
    return client.closeSession(request, options);
  }

  @Override
  public void close() throws Exception {
    AutoCloseables.close(client);
  }

  /** Class to encapsulate Flight SQL bulk ingest request options. * */
  public static class ExecuteIngestOptions {
    private final String table;
    private final TableDefinitionOptions tableDefinitionOptions;
    private final boolean useTemporaryTable;
    private final String catalog;
    private final String schema;
    private final Map<String, String> options;

    /**
     * Constructor.
     *
     * @param table The table to load data into.
     * @param tableDefinitionOptions The behavior for handling the table definition.
     * @param catalog The catalog of the destination table to load data into. If null, a
     *     backend-specific default may be used.
     * @param schema The schema of the destination table to load data into. If null, a
     *     backend-specific default may be used.
     * @param options Backend-specific options. Can be null if there are no options to be set.
     */
    public ExecuteIngestOptions(
        String table,
        TableDefinitionOptions tableDefinitionOptions,
        String catalog,
        String schema,
        Map<String, String> options) {
      this(table, tableDefinitionOptions, false, catalog, schema, options);
    }

    /**
     * Constructor.
     *
     * @param table The table to load data into.
     * @param tableDefinitionOptions The behavior for handling the table definition.
     * @param useTemporaryTable Use a temporary table for bulk ingestion. Temporary table may get
     *     placed in a backend-specific schema and/or catalog and gets dropped at the end of the
     *     session. If backend does not support ingesting using a temporary table or an explicit
     *     choice of schema or catalog is incompatible with the server's namespacing decision, an
     *     error is returned as part of {@link #executeIngest} request.
     * @param catalog The catalog of the destination table to load data into. If null, a
     *     backend-specific default may be used.
     * @param schema The schema of the destination table to load data into. If null, a
     *     backend-specific default may be used.
     * @param options Backend-specific options. Can be null if there are no options to be set.
     */
    public ExecuteIngestOptions(
        String table,
        TableDefinitionOptions tableDefinitionOptions,
        boolean useTemporaryTable,
        String catalog,
        String schema,
        Map<String, String> options) {
      this.table = table;
      this.tableDefinitionOptions = tableDefinitionOptions;
      this.useTemporaryTable = useTemporaryTable;
      this.catalog = catalog;
      this.schema = schema;
      this.options = options;
    }

    protected void updateCommandBuilder(CommandStatementIngest.Builder builder) {
      builder.setTable(table);
      builder.setTableDefinitionOptions(tableDefinitionOptions);
      builder.setTemporary(useTemporaryTable);
      if (!isNull(catalog)) {
        builder.setCatalog(catalog);
      }
      if (!isNull(schema)) {
        builder.setSchema(schema);
      }
      if (!isNull(options)) {
        builder.putAllOptions(options);
      }
    }
  }

  /** Helper class to encapsulate Flight SQL prepared statement logic. */
  public static class PreparedStatement implements AutoCloseable {
    private final FlightClient client;
    private final ActionCreatePreparedStatementResult preparedStatementResult;
    private VectorSchemaRoot parameterBindingRoot;
    private boolean isClosed;
    private Schema resultSetSchema;
    private Schema parameterSchema;

    PreparedStatement(FlightClient client, Action action, CallOption... options) {
      this.client = client;

      final Iterator<Result> preparedStatementResults = client.doAction(action, options);
      preparedStatementResult =
          FlightSqlUtils.unpackAndParseOrThrow(
              preparedStatementResults.next().getBody(), ActionCreatePreparedStatementResult.class);
      isClosed = false;
    }

    /**
     * Set the {@link #parameterBindingRoot} containing the parameter binding from a {@link
     * PreparedStatement} operation.
     *
     * @param parameterBindingRoot a {@code VectorSchemaRoot} object containing the values to be
     *     used in the {@code PreparedStatement} setters.
     */
    public void setParameters(final VectorSchemaRoot parameterBindingRoot) {
      if (parameterBindingRoot == this.parameterBindingRoot) {
        // Nothing to do if we're attempting to set the same parameters again.
        return;
      }
      clearParameters();
      this.parameterBindingRoot = parameterBindingRoot;
    }

    /**
     * Closes the {@link #parameterBindingRoot}, which contains the parameter binding from a {@link
     * PreparedStatement} operation, releasing its resources.
     */
    public void clearParameters() {
      if (parameterBindingRoot != null) {
        parameterBindingRoot.close();
      }
    }

    /**
     * Returns the Schema of the resultset.
     *
     * @return the Schema of the resultset.
     */
    public Schema getResultSetSchema() {
      if (resultSetSchema == null) {
        final ByteString bytes = preparedStatementResult.getDatasetSchema();
        resultSetSchema = deserializeSchema(bytes);
      }
      return resultSetSchema;
    }

    /**
     * Returns the Schema of the parameters.
     *
     * @return the Schema of the parameters.
     */
    public Schema getParameterSchema() {
      if (parameterSchema == null) {
        final ByteString bytes = preparedStatementResult.getParameterSchema();
        parameterSchema = deserializeSchema(bytes);
      }
      return parameterSchema;
    }

    /** Get the schema of the result set (should be identical to {@link #getResultSetSchema()}). */
    public SchemaResult fetchSchema(CallOption... options) {
      checkOpen();

      final FlightDescriptor descriptor =
          FlightDescriptor.command(
              Any.pack(
                      CommandPreparedStatementQuery.newBuilder()
                          .setPreparedStatementHandle(
                              preparedStatementResult.getPreparedStatementHandle())
                          .build())
                  .toByteArray());
      return client.getSchema(descriptor, options);
    }

    private Schema deserializeSchema(final ByteString bytes) {
      try {
        return bytes.isEmpty()
            ? new Schema(Collections.emptyList())
            : MessageSerializer.deserializeSchema(
                new ReadChannel(
                    Channels.newChannel(new ByteArrayInputStream(bytes.toByteArray()))));
      } catch (final IOException e) {
        throw new RuntimeException("Failed to deserialize schema", e);
      }
    }

    /**
     * Executes the prepared statement query on the server.
     *
     * @param options RPC-layer hints for this call.
     * @return a FlightInfo object representing the stream(s) to fetch.
     */
    public FlightInfo execute(final CallOption... options) {
      checkOpen();

      FlightDescriptor descriptor =
          FlightDescriptor.command(
              Any.pack(
                      CommandPreparedStatementQuery.newBuilder()
                          .setPreparedStatementHandle(
                              preparedStatementResult.getPreparedStatementHandle())
                          .build())
                  .toByteArray());

      if (parameterBindingRoot != null && parameterBindingRoot.getRowCount() > 0) {
        try (final SyncPutListener putListener = putParameters(descriptor, options)) {
          if (getParameterSchema().getFields().size() > 0
              && parameterBindingRoot != null
              && parameterBindingRoot.getRowCount() > 0) {
            final PutResult read = putListener.read();
            if (read != null) {
              try (final ArrowBuf metadata = read.getApplicationMetadata()) {
                final FlightSql.DoPutPreparedStatementResult doPutPreparedStatementResult =
                    FlightSql.DoPutPreparedStatementResult.parseFrom(metadata.nioBuffer());
                descriptor =
                    FlightDescriptor.command(
                        Any.pack(
                                CommandPreparedStatementQuery.newBuilder()
                                    .setPreparedStatementHandle(
                                        doPutPreparedStatementResult.getPreparedStatementHandle())
                                    .build())
                            .toByteArray());
              }
            }
          }
        } catch (final InterruptedException e) {
          throw CallStatus.CANCELLED.withCause(e).toRuntimeException();
        } catch (final ExecutionException e) {
          throw CallStatus.CANCELLED.withCause(e.getCause()).toRuntimeException();
        } catch (final InvalidProtocolBufferException e) {
          throw CallStatus.INTERNAL.withCause(e).toRuntimeException();
        }
      }

      return client.getInfo(descriptor, options);
    }

    private SyncPutListener putParameters(FlightDescriptor descriptor, CallOption... options) {
      final SyncPutListener putListener = new SyncPutListener();

      FlightClient.ClientStreamListener listener =
          client.startPut(descriptor, parameterBindingRoot, putListener, options);

      listener.putNext();
      listener.completed();
      listener.getResult();

      return putListener;
    }

    /**
     * Checks whether this client is open.
     *
     * @throws IllegalStateException if client is closed.
     */
    protected final void checkOpen() {
      Preconditions.checkState(!isClosed, "Statement closed");
    }

    /**
     * Executes the prepared statement update on the server.
     *
     * @param options RPC-layer hints for this call.
     * @return the count of updated records
     */
    public long executeUpdate(final CallOption... options) {
      checkOpen();
      final FlightDescriptor descriptor =
          FlightDescriptor.command(
              Any.pack(
                      CommandPreparedStatementUpdate.newBuilder()
                          .setPreparedStatementHandle(
                              preparedStatementResult.getPreparedStatementHandle())
                          .build())
                  .toByteArray());
      setParameters(parameterBindingRoot == null ? VectorSchemaRoot.of() : parameterBindingRoot);
      SyncPutListener putListener = putParameters(descriptor, options);

      try {
        final PutResult read = putListener.read();
        try (final ArrowBuf metadata = read.getApplicationMetadata()) {
          final DoPutUpdateResult doPutUpdateResult =
              DoPutUpdateResult.parseFrom(metadata.nioBuffer());
          return doPutUpdateResult.getRecordCount();
        }
      } catch (final InterruptedException e) {
        throw CallStatus.CANCELLED.withCause(e).toRuntimeException();
      } catch (final ExecutionException e) {
        throw CallStatus.CANCELLED.withCause(e.getCause()).toRuntimeException();
      } catch (final InvalidProtocolBufferException e) {
        throw CallStatus.INTERNAL.withCause(e).toRuntimeException();
      }
    }

    /**
     * Closes the client.
     *
     * @param options RPC-layer hints for this call.
     */
    public void close(final CallOption... options) {
      if (isClosed) {
        return;
      }
      isClosed = true;
      final Action action =
          new Action(
              FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT.getType(),
              Any.pack(
                      ActionClosePreparedStatementRequest.newBuilder()
                          .setPreparedStatementHandle(
                              preparedStatementResult.getPreparedStatementHandle())
                          .build())
                  .toByteArray());
      final Iterator<Result> closePreparedStatementResults = client.doAction(action, options);
      closePreparedStatementResults.forEachRemaining(result -> {});
      clearParameters();
    }

    @Override
    public void close() {
      close(new CallOption[0]);
    }

    /**
     * Returns if the prepared statement is already closed.
     *
     * @return true if the prepared statement is already closed.
     */
    public boolean isClosed() {
      return isClosed;
    }
  }

  /** A handle for an active savepoint. */
  public static class Savepoint {
    private final byte[] transactionId;

    public Savepoint(byte[] transactionId) {
      this.transactionId = transactionId;
    }

    public byte[] getSavepointId() {
      return transactionId;
    }
  }

  /** A handle for an active transaction. */
  public static class Transaction {
    private final byte[] transactionId;

    public Transaction(byte[] transactionId) {
      this.transactionId = transactionId;
    }

    public byte[] getTransactionId() {
      return transactionId;
    }
  }

  /** A wrapper around a Substrait plan and a Substrait version. */
  public static final class SubstraitPlan {
    private final byte[] plan;
    private final String version;

    public SubstraitPlan(byte[] plan, String version) {
      this.plan = Preconditions.checkNotNull(plan);
      this.version = Preconditions.checkNotNull(version);
    }

    public byte[] getPlan() {
      return plan;
    }

    public String getVersion() {
      return version;
    }

    @Override
    public boolean equals(Object o) {
      if (this == o) {
        return true;
      }
      if (o == null || getClass() != o.getClass()) {
        return false;
      }

      SubstraitPlan that = (SubstraitPlan) o;

      if (!Arrays.equals(getPlan(), that.getPlan())) {
        return false;
      }
      return getVersion().equals(that.getVersion());
    }

    @Override
    public int hashCode() {
      int result = Arrays.hashCode(getPlan());
      result = 31 * result + getVersion().hashCode();
      return result;
    }

    @Override
    public String toString() {
      return "SubstraitPlan{"
          + "plan="
          + Arrays.toString(plan)
          + ", version='"
          + version
          + '\''
          + '}';
    }
  }
}
