GH-942: Fix JDBC Connection.setCatalog() (#943)

## What's Changed

Connection.setCatalog() is not silently ignored anymore (through the
default implementation in Calcite) but instead it updates the catalog
session option in the same way as during the initial connection.

Closes #942.

Author: eickler

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

@@ -79,7 +79,8 @@ static Signature newSignature(final String sql, Schema resultSetSchema, Schema p
   public void closeStatement(final StatementHandle statementHandle) {
     PreparedStatement preparedStatement =
         statementHandlePreparedStatementMap.remove(new StatementHandleKey(statementHandle));
-    // Testing if the prepared statement was created because the statement can be not created until
+    // Testing if the prepared statement was created because the statement can be
+    // not created until
     // this moment
     if (preparedStatement != null) {
       preparedStatement.close();
@@ -224,7 +225,8 @@ public ExecuteResult prepareAndExecute(
           MetaResultSet.create(handle.connectionId, handle.id, false, handle.signature, null);
       return new ExecuteResult(Collections.singletonList(metaResultSet));
     } catch (SQLTimeoutException e) {
-      // So far AvaticaStatement(executeInternal) only handles NoSuchStatement and Runtime
+      // So far AvaticaStatement(executeInternal) only handles NoSuchStatement and
+      // Runtime
       // Exceptions.
       throw new RuntimeException(e);
     } catch (SQLException e) {
@@ -253,6 +255,20 @@ public boolean syncResults(
     return false;
   }
 
+  @Override
+  public ConnectionProperties connectionSync(ConnectionHandle ch, ConnectionProperties connProps) {
+    final ConnectionProperties result = super.connectionSync(ch, connProps);
+    final String newCatalog = this.connProps.getCatalog();
+    if (newCatalog != null) {
+      try {
+        ((ArrowFlightConnection) connection).getClientHandler().setCatalog(newCatalog);
+      } catch (SQLException e) {
+        throw new RuntimeException(e);
+      }
+    }
+    return result;
+  }
+
   void setDefaultConnectionProperties() {
     // TODO Double-check this.
     connProps
@@ -268,7 +284,8 @@ PreparedStatement getPreparedStatement(StatementHandle statementHandle) {
     return statementHandlePreparedStatementMap.get(new StatementHandleKey(statementHandle));
   }
 
-  // Helper used to look up prepared statement instances later. Avatica doesn't give us the
+  // Helper used to look up prepared statement instances later. Avatica doesn't
+  // give us the
   // signature in
   // an UPDATE code path so we can't directly use StatementHandle as a map key.
   private static final class StatementHandleKey {

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

@@ -47,7 +47,6 @@
 import org.apache.arrow.flight.FlightStatusCode;
 import org.apache.arrow.flight.Location;
 import org.apache.arrow.flight.LocationSchemes;
-import org.apache.arrow.flight.SessionOptionValue;
 import org.apache.arrow.flight.SessionOptionValueFactory;
 import org.apache.arrow.flight.SetSessionOptionsRequest;
 import org.apache.arrow.flight.SetSessionOptionsResult;
@@ -147,20 +146,26 @@ public List<CloseableEndpointStreamPair> getStreams(final FlightInfo flightInfo)
     try {
       for (FlightEndpoint endpoint : flightInfo.getEndpoints()) {
         if (endpoint.getLocations().isEmpty()) {
-          // Create a stream using the current client only and do not close the client at the end.
+          // Create a stream using the current client only and do not close the client at
+          // the end.
           endpoints.add(
               new CloseableEndpointStreamPair(
                   sqlClient.getStream(endpoint.getTicket(), getOptions()), null));
         } else {
           // Clone the builder and then set the new endpoint on it.
 
-          // GH-38574: Currently a new FlightClient will be made for each partition that returns a
-          // non-empty Location then disposed of. It may be better to cache clients because a server
-          // may report the same Locations. It would also be good to identify when the reported
+          // GH-38574: Currently a new FlightClient will be made for each partition that
+          // returns a
+          // non-empty Location then disposed of. It may be better to cache clients
+          // because a server
+          // may report the same Locations. It would also be good to identify when the
+          // reported
           // location
-          // is the same as the original connection's Location and skip creating a FlightClient in
+          // is the same as the original connection's Location and skip creating a
+          // FlightClient in
           // that scenario.
-          // Also copy the cache to the client so we can share a cache. Cache needs to cache
+          // Also copy the cache to the client so we can share a cache. Cache needs to
+          // cache
           // negative attempts too.
           List<Exception> exceptions = new ArrayList<>();
           CloseableEndpointStreamPair stream = null;
@@ -337,7 +342,8 @@ private boolean isBenignCloseException(FlightRuntimeException fre) {
    */
   private void logSuppressedCloseException(
       FlightRuntimeException fre, String operationDescription) {
-    // ARROW-17785 and GH-863: suppress exceptions caused by flaky gRPC layer during shutdown
+    // ARROW-17785 and GH-863: suppress exceptions caused by flaky gRPC layer during
+    // shutdown
     LOGGER.debug("Suppressed error {}", operationDescription, fre);
   }
 
@@ -388,25 +394,40 @@ public interface PreparedStatement extends AutoCloseable {
   /** A connection is created with catalog set as a session option. */
   private void setSetCatalogInSessionIfPresent() {
     if (catalog.isPresent()) {
-      final SetSessionOptionsRequest setSessionOptionRequest =
-          new SetSessionOptionsRequest(
-              ImmutableMap.<String, SessionOptionValue>builder()
-                  .put(CATALOG, SessionOptionValueFactory.makeSessionOptionValue(catalog.get()))
-                  .build());
-      final SetSessionOptionsResult result =
-          sqlClient.setSessionOptions(setSessionOptionRequest, getOptions());
+      try {
+        setCatalog(catalog.get());
+      } catch (SQLException e) {
+        throw CallStatus.INVALID_ARGUMENT
+            .withDescription(e.getMessage())
+            .withCause(e)
+            .toRuntimeException();
+      }
+    }
+  }
 
+  /**
+   * Sets the catalog for the current session.
+   *
+   * @param catalog the catalog to set.
+   * @throws SQLException if an error occurs while setting the catalog.
+   */
+  public void setCatalog(final String catalog) throws SQLException {
+    final SetSessionOptionsRequest request =
+        new SetSessionOptionsRequest(
+            ImmutableMap.of(CATALOG, SessionOptionValueFactory.makeSessionOptionValue(catalog)));
+    try {
+      final SetSessionOptionsResult result = sqlClient.setSessionOptions(request, getOptions());
       if (result.hasErrors()) {
-        Map<String, SetSessionOptionsResult.Error> errors = result.getErrors();
-        for (Map.Entry<String, SetSessionOptionsResult.Error> error : errors.entrySet()) {
+        final Map<String, SetSessionOptionsResult.Error> errors = result.getErrors();
+        for (final Map.Entry<String, SetSessionOptionsResult.Error> error : errors.entrySet()) {
           LOGGER.warn(error.toString());
         }
-        throw CallStatus.INVALID_ARGUMENT
-            .withDescription(
-                String.format(
-                    "Cannot set session option for catalog = %s. Check log for details.", catalog))
-            .toRuntimeException();
+        throw new SQLException(
+            String.format(
+                "Cannot set session option for catalog = %s. Check log for details.", catalog));
       }
+    } catch (final FlightRuntimeException e) {
+      throw new SQLException(e);
     }
   }
 
@@ -654,7 +675,8 @@ public static final class Builder {
 
     @VisibleForTesting @Nullable Duration connectTimeout;
 
-    // These two middleware are for internal use within build() and should not be exposed by builder
+    // These two middleware are for internal use within build() and should not be
+    // exposed by builder
     // APIs.
     // Note that these middleware may not necessarily be registered.
     @VisibleForTesting
@@ -980,15 +1002,17 @@ public Location getLocation() {
      * @throws SQLException on error.
      */
     public ArrowFlightSqlClientHandler build() throws SQLException {
-      // Copy middleware so that the build method doesn't change the state of the builder fields
+      // Copy middleware so that the build method doesn't change the state of the
+      // builder fields
       // itself.
       Set<FlightClientMiddleware.Factory> buildTimeMiddlewareFactories =
           new HashSet<>(this.middlewareFactories);
       FlightClient client = null;
       boolean isUsingUserPasswordAuth = username != null && token == null;
 
       try {
-        // Token should take priority since some apps pass in a username/password even when a token
+        // Token should take priority since some apps pass in a username/password even
+        // when a token
         // is provided
         if (isUsingUserPasswordAuth) {
           buildTimeMiddlewareFactories.add(authFactory);
@@ -1047,8 +1071,10 @@ public ArrowFlightSqlClientHandler build() throws SQLException {
                 allocator, channelBuilder.build(), clientBuilder.middleware());
         final ArrayList<CallOption> credentialOptions = new ArrayList<>();
         if (isUsingUserPasswordAuth) {
-          // If the authFactory has already been used for a handshake, use the existing token.
-          // This can occur if the authFactory is being re-used for a new connection spawned for
+          // If the authFactory has already been used for a handshake, use the existing
+          // token.
+          // This can occur if the authFactory is being re-used for a new connection
+          // spawned for
           // getStream().
           if (authFactory.getCredentialCallOption() != null) {
             credentialOptions.add(authFactory.getCredentialCallOption());

flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java

@@ -16,6 +16,7 @@
  */
 package org.apache.arrow.driver.jdbc;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -26,12 +27,15 @@
 import java.sql.Driver;
 import java.sql.DriverManager;
 import java.sql.SQLException;
+import java.util.Map;
 import java.util.Properties;
 import org.apache.arrow.driver.jdbc.authentication.UserPasswordAuthentication;
 import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler;
 import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty;
 import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer;
 import org.apache.arrow.flight.FlightMethod;
+import org.apache.arrow.flight.NoOpSessionOptionValueVisitor;
+import org.apache.arrow.flight.SessionOptionValue;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.util.AutoCloseables;
@@ -614,12 +618,46 @@ public void testJdbcDriverVersionIntegration() throws Exception {
 
       var expectedUserAgent =
           "JDBC Flight SQL Driver " + driverVersion.getDriverVersion().versionString;
-      // Driver appends version to grpc user-agent header. Assert the header starts with the
+      // Driver appends version to grpc user-agent header. Assert the header starts
+      // with the
       // expected
       // value and ignored grpc version.
       assertTrue(
           actualUserAgent.startsWith(expectedUserAgent),
           "Expected: " + expectedUserAgent + " but found: " + actualUserAgent);
     }
   }
+
+  @Test
+  public void testSetCatalogShouldUpdateSessionOptions() throws Exception {
+    final Properties properties = new Properties();
+    properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest);
+    properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest);
+    properties.put("useEncryption", false);
+
+    try (Connection connection =
+        DriverManager.getConnection(
+            "jdbc:arrow-flight-sql://"
+                + FLIGHT_SERVER_TEST_EXTENSION.getHost()
+                + ":"
+                + FLIGHT_SERVER_TEST_EXTENSION.getPort(),
+            properties)) {
+      final String catalog = "new_catalog";
+      connection.setCatalog(catalog);
+
+      final Map<String, SessionOptionValue> options = PRODUCER.getSessionOptions();
+      assertTrue(options.containsKey("catalog"));
+      String actualCatalog =
+          options
+              .get("catalog")
+              .acceptVisitor(
+                  new NoOpSessionOptionValueVisitor<String>() {
+                    @Override
+                    public String visit(String value) {
+                      return value;
+                    }
+                  });
+      assertEquals(catalog, actualCatalog);
+    }
+  }
 }

flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java

@@ -52,6 +52,9 @@
 import org.apache.arrow.flight.PutResult;
 import org.apache.arrow.flight.Result;
 import org.apache.arrow.flight.SchemaResult;
+import org.apache.arrow.flight.SessionOptionValue;
+import org.apache.arrow.flight.SetSessionOptionsRequest;
+import org.apache.arrow.flight.SetSessionOptionsResult;
 import org.apache.arrow.flight.Ticket;
 import org.apache.arrow.flight.sql.FlightSqlProducer;
 import org.apache.arrow.flight.sql.SqlInfoBuilder;
@@ -664,6 +667,22 @@ public SqlInfoBuilder getSqlInfoBuilder() {
     return sqlInfoBuilder;
   }
 
+  private final Map<String, SessionOptionValue> sessionOptions = new HashMap<>();
+
+  @Override
+  public void setSessionOptions(
+      final SetSessionOptionsRequest request,
+      final CallContext context,
+      final StreamListener<SetSessionOptionsResult> listener) {
+    sessionOptions.putAll(request.getSessionOptions());
+    listener.onNext(new SetSessionOptionsResult(Collections.emptyMap()));
+    listener.onCompleted();
+  }
+
+  public Map<String, SessionOptionValue> getSessionOptions() {
+    return sessionOptions;
+  }
+
   private static final class TicketConversionUtils {
     private TicketConversionUtils() {
       // Prevent instantiation.