BDX-645: prototype PollInfo support for JDBC

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java

@@ -126,6 +126,7 @@ private static ArrowFlightSqlClientHandler createNewClientHandler(
           .withCatalog(config.getCatalog())
           .withClientCache(config.useClientCache() ? new FlightClientCache() : null)
           .withConnectTimeout(config.getConnectTimeout())
+          .withPollInfo(config.usePollInfo())
           .withDriverVersion(driverVersion)
           .withOAuthConfiguration(config.getOauthConfiguration())
           .build();

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java

@@ -21,13 +21,20 @@
 import java.sql.ResultSet;
 import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
+import java.sql.SQLTimeoutException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
 import java.util.TimeZone;
 import java.util.concurrent.TimeUnit;
 import org.apache.arrow.driver.jdbc.client.CloseableEndpointStreamPair;
+import org.apache.arrow.driver.jdbc.client.PollInfoOperation;
 import org.apache.arrow.driver.jdbc.utils.FlightEndpointDataQueue;
 import org.apache.arrow.driver.jdbc.utils.VectorSchemaRootTransformer;
+import org.apache.arrow.flight.FlightEndpoint;
 import org.apache.arrow.flight.FlightInfo;
+import org.apache.arrow.flight.FlightRuntimeException;
+import org.apache.arrow.flight.FlightStatusCode;
 import org.apache.arrow.flight.FlightStream;
 import org.apache.arrow.util.AutoCloseables;
 import org.apache.arrow.vector.VectorSchemaRoot;
@@ -46,7 +53,9 @@ public final class ArrowFlightJdbcFlightStreamResultSet
     extends ArrowFlightJdbcVectorSchemaRootResultSet {
 
   private final ArrowFlightConnection connection;
-  private final FlightInfo flightInfo;
+  private FlightInfo flightInfo;
+  private final PollInfoOperation pollInfoOperation;
+  private int consumedEndpointCount;
   private CloseableEndpointStreamPair currentEndpointData;
   private FlightEndpointDataQueue flightEndpointDataQueue;
 
@@ -67,7 +76,21 @@ public final class ArrowFlightJdbcFlightStreamResultSet
       throws SQLException {
     super(statement, state, signature, resultSetMetaData, timeZone, firstFrame);
     this.connection = (ArrowFlightConnection) statement.connection;
-    this.flightInfo = ((ArrowFlightInfoStatement) statement).executeFlightInfoQuery();
+    try {
+      this.flightInfo = ((ArrowFlightInfoStatement) statement).executeFlightInfoQuery();
+      this.pollInfoOperation = activePollInfoOperation(statement);
+    } catch (FlightRuntimeException e) {
+      if (e.status().code() != FlightStatusCode.TIMED_OUT) {
+        throw e;
+      }
+      final SQLTimeoutException jdbcTimeout =
+          new SQLTimeoutException(
+              String.format(
+                  "Query timed out after %d %s",
+                  statement.getQueryTimeout(), TimeUnit.SECONDS));
+      jdbcTimeout.initCause(e);
+      throw jdbcTimeout;
+    }
   }
 
   /** Private constructor for fromFlightInfo. */
@@ -83,6 +106,7 @@ private ArrowFlightJdbcFlightStreamResultSet(
     super(null, state, signature, resultSetMetaData, timeZone, firstFrame);
     this.connection = connection;
     this.flightInfo = flightInfo;
+    this.pollInfoOperation = null;
     this.id = connection.getNewMetadataResultSetId(this);
   }
 
@@ -143,7 +167,7 @@ protected AvaticaResultSet execute() throws SQLException {
 
   private void populateData() throws SQLException {
     loadNewQueue();
-    flightEndpointDataQueue.enqueue(connection.getClientHandler().getStreams(flightInfo));
+    enqueueNewEndpoints(flightInfo);
     loadNewFlightStream();
 
     // Ownership of the root will be passed onto the cursor.
@@ -176,7 +200,9 @@ public byte[] getAppMetadata() {
   @Override
   public boolean next() throws SQLException {
     if (currentVectorSchemaRoot == null) {
-      return false;
+      if (!loadNextPublishedEndpoint()) {
+        return false;
+      }
     }
     while (true) {
       final boolean hasNext = super.next();
@@ -209,6 +235,10 @@ public boolean next() throws SQLException {
         continue;
       }
 
+      if (loadNextPublishedEndpoint()) {
+        continue;
+      }
+
       if (statement != null && statement.isCloseOnCompletion()) {
         statement.close();
       }
@@ -219,6 +249,10 @@ public boolean next() throws SQLException {
 
   @Override
   protected void cancel() {
+    if (pollInfoOperation != null && !pollInfoOperation.isComplete()) {
+      pollInfoOperation.cancel();
+    }
+    finishPollInfoOperation();
     super.cancel();
     final CloseableEndpointStreamPair currentEndpoint = this.currentEndpointData;
     if (currentEndpoint != null) {
@@ -253,19 +287,134 @@ public synchronized void close() {
     } catch (final Exception e) {
       throw new RuntimeException(e);
     } finally {
+      if (pollInfoOperation != null && !pollInfoOperation.isComplete()) {
+        pollInfoOperation.cancel();
+      }
+      finishPollInfoOperation();
       super.close();
     }
   }
 
   private CloseableEndpointStreamPair getNextEndpointStream(final boolean canTimeout)
       throws SQLException {
     if (canTimeout) {
-      final int statementTimeout = statement != null ? statement.getQueryTimeout() : 0;
-      return statementTimeout != 0
-          ? flightEndpointDataQueue.next(statementTimeout, TimeUnit.SECONDS)
-          : flightEndpointDataQueue.next();
+      final long remainingTimeoutNanos = remainingQueryTimeoutNanos();
+      if (remainingTimeoutNanos != Long.MAX_VALUE) {
+        if (remainingTimeoutNanos <= 0) {
+          throw new SQLTimeoutException("Query timed out before retrieving its first endpoint");
+        }
+        try {
+          return flightEndpointDataQueue.next(remainingTimeoutNanos, TimeUnit.NANOSECONDS);
+        } catch (SQLTimeoutException e) {
+          if (statement != null && e.getMessage().startsWith("Query timed out after")) {
+            final SQLTimeoutException jdbcTimeout =
+                new SQLTimeoutException(
+                    String.format(
+                        "Query timed out after %d %s",
+                        statement.getQueryTimeout(), TimeUnit.SECONDS));
+            jdbcTimeout.initCause(e);
+            throw jdbcTimeout;
+          }
+          throw e;
+        }
+      }
     } else {
       return flightEndpointDataQueue.next();
     }
+    return flightEndpointDataQueue.next();
+  }
+
+  private long remainingQueryTimeoutNanos() throws SQLException {
+    if (statement instanceof ArrowFlightStatement) {
+      return ((ArrowFlightStatement) statement).remainingQueryTimeoutNanos();
+    }
+    if (statement instanceof ArrowFlightPreparedStatement) {
+      return ((ArrowFlightPreparedStatement) statement).remainingQueryTimeoutNanos();
+    }
+    final int statementTimeout = statement != null ? statement.getQueryTimeout() : 0;
+    return statementTimeout > 0 ? TimeUnit.SECONDS.toNanos(statementTimeout) : Long.MAX_VALUE;
+  }
+
+  private void enqueueNewEndpoints(final FlightInfo updatedFlightInfo) throws SQLException {
+    final int updatedEndpointCount = updatedFlightInfo.getEndpoints().size();
+    if (updatedEndpointCount < consumedEndpointCount) {
+      throw new SQLException("PollInfo removed previously published endpoints");
+    }
+    final Schema updatedSchema = updatedFlightInfo.getSchemaOptional().orElse(schema);
+    if (schema != null && updatedSchema != null && !schema.equals(updatedSchema)) {
+      throw new SQLException("PollInfo changed the result schema");
+    }
+    if (updatedEndpointCount == consumedEndpointCount) {
+      flightInfo = updatedFlightInfo;
+      return;
+    }
+    final List<FlightEndpoint> appendedEndpoints =
+        new ArrayList<>(
+            updatedFlightInfo
+                .getEndpoints()
+                .subList(consumedEndpointCount, updatedEndpointCount));
+    final FlightInfo appendedInfo =
+        new FlightInfo(
+            updatedSchema,
+            updatedFlightInfo.getDescriptor(),
+            appendedEndpoints,
+            updatedFlightInfo.getBytes(),
+            updatedFlightInfo.getRecords());
+    flightEndpointDataQueue.enqueue(connection.getClientHandler().getStreams(appendedInfo));
+    consumedEndpointCount = updatedEndpointCount;
+    flightInfo = updatedFlightInfo;
+  }
+
+  private boolean loadNextPublishedEndpoint() throws SQLException {
+    if (pollInfoOperation == null || !pollInfoOperation.hasContinuation()) {
+      finishPollInfoOperation();
+      return false;
+    }
+    try {
+      enqueueNewEndpoints(pollInfoOperation.pollNextAvailable());
+      currentEndpointData = getNextEndpointStream(false);
+      if (currentEndpointData != null) {
+        populateDataForCurrentFlightStream();
+        return true;
+      }
+      finishPollInfoOperation();
+      return false;
+    } catch (FlightRuntimeException e) {
+      if (e.status().code() == FlightStatusCode.CANCELLED
+          || e.status().code() == FlightStatusCode.TIMED_OUT) {
+        pollInfoOperation.cancel();
+      } else {
+        pollInfoOperation.terminate();
+      }
+      finishPollInfoOperation();
+      if (e.status().code() == FlightStatusCode.TIMED_OUT) {
+        final SQLTimeoutException timeout =
+            new SQLTimeoutException(
+                String.format(
+                    "Query timed out after %d %s",
+                    statement.getQueryTimeout(), TimeUnit.SECONDS));
+        timeout.initCause(e);
+        throw timeout;
+      }
+      throw new SQLException("Continuation PollFlightInfo failed", e);
+    }
+  }
+
+  private static PollInfoOperation activePollInfoOperation(final AvaticaStatement statement) {
+    if (statement instanceof ArrowFlightStatement) {
+      return ((ArrowFlightStatement) statement).activePollInfoOperation();
+    }
+    if (statement instanceof ArrowFlightPreparedStatement) {
+      return ((ArrowFlightPreparedStatement) statement).activePollInfoOperation();
+    }
+    return null;
+  }
+
+  private void finishPollInfoOperation() {
+    if (statement instanceof ArrowFlightStatement) {
+      ((ArrowFlightStatement) statement).finishPollInfoOperation(pollInfoOperation);
+    } else if (statement instanceof ArrowFlightPreparedStatement) {
+      ((ArrowFlightPreparedStatement) statement).finishPollInfoOperation(pollInfoOperation);
+    }
   }
 }

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java

@@ -181,9 +181,12 @@ public Frame fetch(
         String.format("%s does not use frames.", this), AvaticaConnection.HELPER.unsupported());
   }
 
-  private PreparedStatement prepareForHandle(final String query, StatementHandle handle) {
+  private PreparedStatement prepareForHandle(
+      final String query, StatementHandle handle, final boolean directExecution) {
     final PreparedStatement preparedStatement =
-        ((ArrowFlightConnection) connection).getClientHandler().prepare(query);
+        directExecution
+            ? ((ArrowFlightConnection) connection).getClientHandler().prepareDirect(query)
+            : ((ArrowFlightConnection) connection).getClientHandler().prepare(query);
     handle.signature =
         newSignature(
             query,
@@ -198,7 +201,7 @@ private PreparedStatement prepareForHandle(final String query, StatementHandle h
   public StatementHandle prepare(
       final ConnectionHandle connectionHandle, final String query, final long maxRowCount) {
     final StatementHandle handle = super.createStatement(connectionHandle);
-    prepareForHandle(query, handle);
+    prepareForHandle(query, handle, false);
     return handle;
   }
 
@@ -222,7 +225,7 @@ public ExecuteResult prepareAndExecute(
       final PrepareCallback callback)
       throws NoSuchStatementException {
     try {
-      PreparedStatement preparedStatement = prepareForHandle(query, handle);
+      PreparedStatement preparedStatement = prepareForHandle(query, handle, true);
       final StatementType statementType = preparedStatement.getType();
 
       final long updateCount =

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java

@@ -18,7 +18,9 @@
 
 import java.sql.PreparedStatement;
 import java.sql.SQLException;
+import java.util.concurrent.atomic.AtomicReference;
 import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler;
+import org.apache.arrow.driver.jdbc.client.PollInfoOperation;
 import org.apache.arrow.flight.FlightInfo;
 import org.apache.arrow.util.Preconditions;
 import org.apache.calcite.avatica.AvaticaPreparedStatement;
@@ -30,6 +32,8 @@ public class ArrowFlightPreparedStatement extends AvaticaPreparedStatement
     implements ArrowFlightInfoStatement {
 
   private final ArrowFlightSqlClientHandler.PreparedStatement preparedStatement;
+  private final AtomicReference<PollInfoOperation> activeOperation = new AtomicReference<>();
+  private volatile PollInfoOperation lastOperation;
 
   private ArrowFlightPreparedStatement(
       final ArrowFlightConnection connection,
@@ -76,6 +80,44 @@ public synchronized void close() throws SQLException {
 
   @Override
   public FlightInfo executeFlightInfoQuery() throws SQLException {
-    return preparedStatement.executeQuery();
+    final PollInfoOperation operation = new PollInfoOperation(getQueryTimeout(), true);
+    activeOperation.set(operation);
+    lastOperation = operation;
+    try {
+      final FlightInfo flightInfo = preparedStatement.executeQuery(operation);
+      if (!operation.hasContinuation()) {
+        finishPollInfoOperation(operation);
+      }
+      return flightInfo;
+    } catch (RuntimeException | SQLException e) {
+      finishPollInfoOperation(operation);
+      throw e;
+    }
+  }
+
+  @Override
+  public void cancel() throws SQLException {
+    final PollInfoOperation operation = activeOperation.get();
+    if (operation != null) {
+      operation.cancel();
+    }
+    super.cancel();
+  }
+
+  long remainingQueryTimeoutNanos() {
+    final PollInfoOperation operation = lastOperation;
+    return operation != null && operation.hasDeadline()
+        ? operation.remainingTimeoutNanos()
+        : Long.MAX_VALUE;
+  }
+
+  PollInfoOperation activePollInfoOperation() {
+    return activeOperation.get();
+  }
+
+  void finishPollInfoOperation(final PollInfoOperation operation) {
+    if (operation != null && activeOperation.compareAndSet(operation, null)) {
+      operation.close();
+    }
   }
 }