You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by jb...@apache.org on 2026/01/05 06:39:20 UTC
(arrow-java) branch main updated: GH-942: Fix JDBC Connection.setCatalog() (#943)
This is an automated email from the ASF dual-hosted git repository.
jbonofre pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-java.git
The following commit(s) were added to refs/heads/main by this push:
new b9e40fa0a GH-942: Fix JDBC Connection.setCatalog() (#943)
b9e40fa0a is described below
commit b9e40fa0a90143eef36ca53d2f937f2fef494402
Author: André Eickler <79...@users.noreply.github.com>
AuthorDate: Mon Jan 5 07:39:15 2026 +0100
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.
---
.../arrow/driver/jdbc/ArrowFlightMetaImpl.java | 23 ++++++-
.../jdbc/client/ArrowFlightSqlClientHandler.java | 80 ++++++++++++++--------
.../apache/arrow/driver/jdbc/ConnectionTest.java | 40 ++++++++++-
.../driver/jdbc/utils/MockFlightSqlProducer.java | 19 +++++
4 files changed, 131 insertions(+), 31 deletions(-)
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java
index 21cc3e431..64529b50c 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java
@@ -79,7 +79,8 @@ public class ArrowFlightMetaImpl extends MetaImpl {
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 class ArrowFlightMetaImpl extends MetaImpl {
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 class ArrowFlightMetaImpl extends MetaImpl {
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 @@ public class ArrowFlightMetaImpl extends MetaImpl {
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 {
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java
index 5dc7e0e2e..666996cd9 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java
+++ b/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.FlightRuntimeException;
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 final class ArrowFlightSqlClientHandler implements AutoCloseable {
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 @@ public final class ArrowFlightSqlClientHandler implements AutoCloseable {
*/
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 final class ArrowFlightSqlClientHandler implements 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 final class ArrowFlightSqlClientHandler implements AutoCloseable {
@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,7 +1002,8 @@ public final class ArrowFlightSqlClientHandler implements AutoCloseable {
* @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);
@@ -988,7 +1011,8 @@ public final class ArrowFlightSqlClientHandler implements AutoCloseable {
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 final class ArrowFlightSqlClientHandler implements AutoCloseable {
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());
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java
index 72e4b222a..46762f331 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java
+++ b/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.Connection;
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,7 +618,8 @@ public class ConnectionTest {
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(
@@ -622,4 +627,37 @@ public class ConnectionTest {
"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);
+ }
+ }
}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java
index a8874c486..45c2a9640 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java
+++ b/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.FlightStream;
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 final class MockFlightSqlProducer implements FlightSqlProducer {
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.