org.apache.arrow.adbc.driver.flightsql.FlightSqlClientWithCallOptions Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of adbc-driver-flight-sql Show documentation
An ADBC driver wrapping Flight SQL.
The newest version!
/*
 * 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.adbc.driver.flightsql;

import static org.apache.arrow.flight.sql.FlightSqlClient.PreparedStatement;
import static org.apache.arrow.flight.sql.FlightSqlClient.Savepoint;
import static org.apache.arrow.flight.sql.FlightSqlClient.SubstraitPlan;
import static org.apache.arrow.flight.sql.FlightSqlClient.Transaction;

import java.util.List;
import org.apache.arrow.flight.CallOption;
import org.apache.arrow.flight.CancelFlightInfoRequest;
import org.apache.arrow.flight.CancelFlightInfoResult;
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.RenewFlightEndpointRequest;
import org.apache.arrow.flight.SchemaResult;
import org.apache.arrow.flight.Ticket;
import org.apache.arrow.flight.sql.CancelResult;
import org.apache.arrow.flight.sql.FlightSqlClient;
import org.apache.arrow.flight.sql.impl.FlightSql;
import org.apache.arrow.flight.sql.util.TableRef;
import org.apache.arrow.util.AutoCloseables;
import org.checkerframework.checker.nullness.qual.Nullable;

/** A wrapper around FlightSqlClient which automatically adds CallOptions to each RPC call. */
public class FlightSqlClientWithCallOptions implements AutoCloseable {
  private final FlightSqlClient client;
  private final CallOption[] connectionOptions;

  public FlightSqlClientWithCallOptions(FlightSqlClient client, CallOption... options) {
    this.client = client;
    this.connectionOptions = options;
  }

  public FlightInfo execute(String query, CallOption... options) {
    return client.execute(query, combine(options));
  }

  public FlightInfo execute(String query, Transaction transaction, CallOption... options) {
    return client.execute(query, transaction, combine(options));
  }

  public FlightInfo executeSubstrait(SubstraitPlan plan, CallOption... options) {
    return client.executeSubstrait(plan, combine(options));
  }

  public FlightInfo executeSubstrait(
      SubstraitPlan plan, Transaction transaction, CallOption... options) {
    return client.executeSubstrait(plan, transaction, combine(options));
  }

  public SchemaResult getExecuteSchema(
      String query, Transaction transaction, CallOption... options) {
    return client.getExecuteSchema(query, transaction, combine(options));
  }

  public SchemaResult getExecuteSchema(String query, CallOption... options) {
    return client.getExecuteSchema(query, combine(options));
  }

  public SchemaResult getExecuteSubstraitSchema(
      SubstraitPlan plan, Transaction transaction, CallOption... options) {
    return client.getExecuteSubstraitSchema(plan, transaction, combine(options));
  }

  public SchemaResult getExecuteSubstraitSchema(
      SubstraitPlan substraitPlan, CallOption... options) {
    return client.getExecuteSubstraitSchema(substraitPlan, combine(options));
  }

  public long executeUpdate(String query, CallOption... options) {
    return client.executeUpdate(query, combine(options));
  }

  public long executeUpdate(String query, Transaction transaction, CallOption... options) {
    return client.executeUpdate(query, transaction, combine(options));
  }

  public long executeSubstraitUpdate(SubstraitPlan plan, CallOption... options) {
    return client.executeSubstraitUpdate(plan, combine(options));
  }

  public long executeSubstraitUpdate(
      SubstraitPlan plan, Transaction transaction, CallOption... options) {
    return client.executeSubstraitUpdate(plan, transaction, combine(options));
  }

  public FlightInfo getCatalogs(CallOption... options) {
    return client.getCatalogs(combine(options));
  }

  public SchemaResult getCatalogsSchema(CallOption... options) {
    return client.getCatalogsSchema(combine(options));
  }

  public FlightInfo getSchemas(
      String catalog, String dbSchemaFilterPattern, CallOption... options) {
    return client.getSchemas(catalog, dbSchemaFilterPattern, combine(options));
  }

  public SchemaResult getSchemasSchema(CallOption... options) {
    return client.getSchemasSchema(combine(options));
  }

  public SchemaResult getSchema(FlightDescriptor descriptor, CallOption... options) {
    return client.getSchema(descriptor, combine(options));
  }

  public FlightStream getStream(Ticket ticket, CallOption... options) {
    return client.getStream(ticket, combine(options));
  }

  public FlightInfo getSqlInfo(FlightSql.SqlInfo... info) {
    return client.getSqlInfo(info, combine());
  }

  public FlightInfo getSqlInfo(FlightSql.SqlInfo[] info, CallOption... options) {
    return client.getSqlInfo(info, combine(options));
  }

  public FlightInfo getSqlInfo(int[] info, CallOption... options) {
    return client.getSqlInfo(info, combine(options));
  }

  public FlightInfo getSqlInfo(Iterable info, CallOption... options) {
    return client.getSqlInfo(info, combine(options));
  }

  public SchemaResult getSqlInfoSchema(CallOption... options) {
    return client.getSqlInfoSchema(combine(options));
  }

  public FlightInfo getXdbcTypeInfo(int dataType, CallOption... options) {
    return client.getXdbcTypeInfo(dataType, combine(options));
  }

  public FlightInfo getXdbcTypeInfo(CallOption... options) {
    return client.getXdbcTypeInfo(combine(options));
  }

  public SchemaResult getXdbcTypeInfoSchema(CallOption... options) {
    return client.getXdbcTypeInfoSchema(combine(options));
  }

  @SuppressWarnings("argument")
  // FlightSqlClient.getTableTyoes() allows for null tableTypes but isn't annotated correctly.
  public FlightInfo getTables(
      String catalog,
      String dbSchemaFilterPattern,
      String tableFilterPattern,
      @Nullable List tableTypes,
      boolean includeSchema,
      CallOption... options) {
    return client.getTables(
        catalog,
        dbSchemaFilterPattern,
        tableFilterPattern,
        tableTypes,
        includeSchema,
        combine(options));
  }

  public SchemaResult getTablesSchema(boolean includeSchema, CallOption... options) {
    return client.getTablesSchema(includeSchema, combine(options));
  }

  public FlightInfo getPrimaryKeys(TableRef tableRef, CallOption... options) {
    return client.getPrimaryKeys(tableRef, combine(options));
  }

  public SchemaResult getPrimaryKeysSchema(CallOption... options) {
    return client.getPrimaryKeysSchema(combine(options));
  }

  public FlightInfo getExportedKeys(TableRef tableRef, CallOption... options) {
    return client.getExportedKeys(tableRef, combine(options));
  }

  public SchemaResult getExportedKeysSchema(CallOption... options) {
    return client.getExportedKeysSchema(combine(options));
  }

  public FlightInfo getImportedKeys(TableRef tableRef, CallOption... options) {
    return client.getImportedKeys(tableRef, combine(options));
  }

  public SchemaResult getImportedKeysSchema(CallOption... options) {
    return client.getImportedKeysSchema(combine(options));
  }

  public FlightInfo getCrossReference(
      TableRef pkTableRef, TableRef fkTableRef, CallOption... options) {
    return client.getCrossReference(pkTableRef, fkTableRef, combine(options));
  }

  public SchemaResult getCrossReferenceSchema(CallOption... options) {
    return client.getCrossReferenceSchema(combine(options));
  }

  public FlightInfo getTableTypes(CallOption... options) {
    return client.getTableTypes(combine(options));
  }

  public SchemaResult getTableTypesSchema(CallOption... options) {
    return client.getTableTypesSchema(combine(options));
  }

  public PreparedStatement prepare(String query, CallOption... options) {
    return client.prepare(query, combine(options));
  }

  public PreparedStatement prepare(
      String query, FlightSqlClient.Transaction transaction, CallOption... options) {
    return client.prepare(query, transaction, combine(options));
  }

  public PreparedStatement prepare(SubstraitPlan plan, CallOption... options) {
    return client.prepare(plan, combine(options));
  }

  public PreparedStatement prepare(
      SubstraitPlan plan, Transaction transaction, CallOption... options) {
    return client.prepare(plan, transaction, combine(options));
  }

  public Transaction beginTransaction(CallOption... options) {
    return client.beginTransaction(combine(options));
  }

  public Savepoint beginSavepoint(Transaction transaction, String name, CallOption... options) {
    return client.beginSavepoint(transaction, name, combine(options));
  }

  public void commit(Transaction transaction, CallOption... options) {
    client.commit(transaction, combine(options));
  }

  public void release(Savepoint savepoint, CallOption... options) {
    client.release(savepoint, combine(options));
  }

  public void rollback(Transaction transaction, CallOption... options) {
    client.rollback(transaction, combine(options));
  }

  public void rollback(Savepoint savepoint, CallOption... options) {
    client.rollback(savepoint, combine(options));
  }

  public CancelFlightInfoResult cancelFlightInfo(
      CancelFlightInfoRequest request, CallOption... options) {
    return client.cancelFlightInfo(request, combine(options));
  }

  public CancelResult cancelQuery(FlightInfo info, CallOption... options) {
    return client.cancelQuery(info, combine(options));
  }

  public FlightEndpoint renewFlightEndpoint(
      RenewFlightEndpointRequest request, CallOption... options) {
    return client.renewFlightEndpoint(request, combine(options));
  }

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

  private CallOption[] combine(CallOption... options) {
    final CallOption[] result = new CallOption[connectionOptions.length + options.length];
    System.arraycopy(connectionOptions, 0, result, 0, connectionOptions.length);
    System.arraycopy(options, 0, result, connectionOptions.length, options.length);
    return result;
  }
}