From f696d46e79e0eff77941b73707c8505dc12cf63a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Thu, 31 Jul 2025 09:59:24 +0100 Subject: [PATCH 1/4] feat: pass down driver version to nettyclientbuild --- .../arrow/flight/grpc/NettyClientBuilder.java | 9 ++++++- .../driver/jdbc/ArrowFlightConnection.java | 9 +++++-- .../client/ArrowFlightSqlClientHandler.java | 25 +++++++++++++++++++ 3 files changed, 40 insertions(+), 3 deletions(-) diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/NettyClientBuilder.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/NettyClientBuilder.java index 42cdaac016..d237812051 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/NettyClientBuilder.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/NettyClientBuilder.java @@ -60,6 +60,7 @@ public class NettyClientBuilder { protected String overrideHostname = null; protected List middleware = new ArrayList<>(); protected boolean verifyServer = true; + protected String userAgent; public NettyClientBuilder() {} @@ -130,6 +131,11 @@ public NettyClientBuilder verifyServer(boolean verifyServer) { return this; } + public NettyClientBuilder userAgent(String userAgent) { + this.userAgent = userAgent; + return this; + } + /** Create the client from this builder. */ public NettyChannelBuilder build() { final NettyChannelBuilder builder; @@ -226,7 +232,8 @@ public NettyChannelBuilder build() { builder .maxTraceEvents(MAX_CHANNEL_TRACE_EVENTS) .maxInboundMessageSize(maxInboundMessageSize) - .maxInboundMetadataSize(maxInboundMessageSize); + .maxInboundMetadataSize(maxInboundMessageSize) + .userAgent(userAgent); return builder; } } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java index 747287ed13..f6f17770f1 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java @@ -32,6 +32,7 @@ import org.apache.arrow.util.Preconditions; import org.apache.calcite.avatica.AvaticaConnection; import org.apache.calcite.avatica.AvaticaFactory; +import org.apache.calcite.avatica.DriverVersion; /** Connection to the Arrow Flight server. */ public final class ArrowFlightConnection extends AvaticaConnection { @@ -86,13 +87,16 @@ static ArrowFlightConnection createNewConnection( throws SQLException { url = replaceSemiColons(url); final ArrowFlightConnectionConfigImpl config = new ArrowFlightConnectionConfigImpl(properties); - final ArrowFlightSqlClientHandler clientHandler = createNewClientHandler(config, allocator); + final ArrowFlightSqlClientHandler clientHandler = + createNewClientHandler(config, allocator, driver.getDriverVersion()); return new ArrowFlightConnection( driver, factory, url, properties, config, allocator, clientHandler); } private static ArrowFlightSqlClientHandler createNewClientHandler( - final ArrowFlightConnectionConfigImpl config, final BufferAllocator allocator) + final ArrowFlightConnectionConfigImpl config, + final BufferAllocator allocator, + final DriverVersion driverVersion) throws SQLException { try { return new ArrowFlightSqlClientHandler.Builder() @@ -116,6 +120,7 @@ private static ArrowFlightSqlClientHandler createNewClientHandler( .withCatalog(config.getCatalog()) .withClientCache(config.useClientCache() ? new FlightClientCache() : null) .withConnectTimeout(config.getConnectTimeout()) + .withDriverVersion(driverVersion) .build(); } catch (final SQLException e) { try { 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 17c2c16ebf..971a3f3f7c 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 @@ -66,6 +66,7 @@ import org.apache.arrow.util.VisibleForTesting; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.calcite.avatica.DriverVersion; import org.apache.calcite.avatica.Meta.StatementType; import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; @@ -548,6 +549,9 @@ public FlightInfo getCrossReference( /** Builder for {@link ArrowFlightSqlClientHandler}. */ public static final class Builder { + @VisibleForTesting static final String USER_AGENT_TEMPLATE = "JDBC Flight SQL Driver %s"; + static final String DEFAULT_VERSION = "(unknown or development build)"; + private final Set middlewareFactories = new HashSet<>(); private final Set options = new HashSet<>(); private String host; @@ -597,6 +601,8 @@ public static final class Builder { @VisibleForTesting ClientCookieMiddleware.Factory cookieFactory = new ClientCookieMiddleware.Factory(); + DriverVersion driverVersion; + public Builder() {} /** @@ -631,6 +637,8 @@ public Builder() {} if (original.retainAuth) { this.authFactory = original.authFactory; } + + this.driverVersion = original.driverVersion; } /** @@ -879,6 +887,17 @@ public Builder withConnectTimeout(Duration connectTimeout) { return this; } + /** + * Sets the driver version for this handler. + * + * @param driverVersion the driver version to set + * @return this builder instance + */ + public Builder withDriverVersion(DriverVersion driverVersion) { + this.driverVersion = driverVersion; + return this; + } + public String getCacheKey() { return getLocation().toString(); } @@ -914,6 +933,12 @@ public ArrowFlightSqlClientHandler build() throws SQLException { final NettyClientBuilder clientBuilder = new NettyClientBuilder(); clientBuilder.allocator(allocator); + String userAgent = String.format(USER_AGENT_TEMPLATE, DEFAULT_VERSION); + if (driverVersion != null && driverVersion.versionString != null) { + userAgent = String.format(USER_AGENT_TEMPLATE, driverVersion.versionString); + } + clientBuilder.userAgent(userAgent); + buildTimeMiddlewareFactories.add(new ClientCookieMiddleware.Factory()); buildTimeMiddlewareFactories.forEach(clientBuilder::intercept); if (useEncryption) { From 17278bd58786981bad961c2881323252aa4fef90 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Thu, 31 Jul 2025 10:01:12 +0100 Subject: [PATCH 2/4] test: assert user-agent is correctly built --- .../ArrowFlightJdbcConnectionCookieTest.java | 4 +- .../arrow/driver/jdbc/ConnectionTest.java | 46 +++++++++++++++++++ .../jdbc/FlightServerTestExtension.java | 33 +++++++++---- ...rrowFlightSqlClientHandlerBuilderTest.java | 1 + 4 files changed, 73 insertions(+), 11 deletions(-) diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionCookieTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionCookieTest.java index 1977b61392..7127c7fc32 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionCookieTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionCookieTest.java @@ -39,11 +39,11 @@ public void testCookies() throws SQLException { Statement statement = connection.createStatement()) { // Expect client didn't receive cookies before any operation - assertNull(FLIGHT_SERVER_TEST_EXTENSION.getMiddlewareCookieFactory().getCookie()); + assertNull(FLIGHT_SERVER_TEST_EXTENSION.getInterceptorFactory().getCookie()); // Run another action for check if the cookies was sent by the server. statement.execute(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD); - assertEquals("k=v", FLIGHT_SERVER_TEST_EXTENSION.getMiddlewareCookieFactory().getCookie()); + assertEquals("k=v", FLIGHT_SERVER_TEST_EXTENSION.getInterceptorFactory().getCookie()); } } } 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 8e872a1167..72e4b222a3 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 @@ -31,6 +31,7 @@ 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.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; @@ -576,4 +577,49 @@ public void testPasswordConnectionPropertyIntegerCorrectCastUrlWithDriverManager assertTrue(connection.isValid(0)); } } + + /** + * Test that the JDBC driver properly integrates driver version into client handler. + * + * @throws Exception on error. + */ + @Test + public void testJdbcDriverVersionIntegration() throws Exception { + final Properties properties = new Properties(); + properties.put( + ArrowFlightConnectionProperty.HOST.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getHost()); + properties.put( + ArrowFlightConnectionProperty.PORT.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getPort()); + properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest); + properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); + properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), false); + + // Create a driver instance and connect + ArrowFlightJdbcDriver driverVersion = new ArrowFlightJdbcDriver(); + + try (Connection connection = + ArrowFlightConnection.createNewConnection( + driverVersion, + new ArrowFlightJdbcFactory(), + "jdbc:arrow-flight-sql://localhost:" + FLIGHT_SERVER_TEST_EXTENSION.getPort(), + properties, + allocator)) { + + assertTrue(connection.isValid(0)); + + var actualUserAgent = + FLIGHT_SERVER_TEST_EXTENSION + .getInterceptorFactory() + .getHeader(FlightMethod.HANDSHAKE, "user-agent"); + + var expectedUserAgent = + "JDBC Flight SQL Driver " + driverVersion.getDriverVersion().versionString; + // 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); + } + } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java index aa586651f5..db0438059f 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java @@ -25,6 +25,8 @@ import java.sql.SQLException; import java.util.ArrayDeque; import java.util.Deque; +import java.util.HashMap; +import java.util.Map; import java.util.Properties; import org.apache.arrow.driver.jdbc.authentication.Authentication; import org.apache.arrow.driver.jdbc.authentication.TokenAuthentication; @@ -33,6 +35,7 @@ import org.apache.arrow.flight.CallHeaders; import org.apache.arrow.flight.CallInfo; import org.apache.arrow.flight.CallStatus; +import org.apache.arrow.flight.FlightMethod; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightServerMiddleware; import org.apache.arrow.flight.Location; @@ -67,7 +70,8 @@ public class FlightServerTestExtension private final CertKeyPair certKeyPair; private final File mTlsCACert; - private final MiddlewareCookie.Factory middlewareCookieFactory = new MiddlewareCookie.Factory(); + private final InterceptorMiddleware.Factory interceptorFactory = + new InterceptorMiddleware.Factory(); private FlightServerTestExtension( final Properties properties, @@ -130,8 +134,8 @@ private void setUseEncryption(boolean useEncryption) { properties.put("useEncryption", useEncryption); } - public MiddlewareCookie.Factory getMiddlewareCookieFactory() { - return middlewareCookieFactory; + public InterceptorMiddleware.Factory getInterceptorFactory() { + return interceptorFactory; } @FunctionalInterface @@ -143,7 +147,7 @@ private FlightServer initiateServer(Location location) throws IOException { FlightServer.Builder builder = FlightServer.builder(allocator, location, producer) .headerAuthenticator(authentication.authenticate()) - .middleware(FlightServerMiddleware.Key.of("KEY"), middlewareCookieFactory); + .middleware(FlightServerMiddleware.Key.of("KEY"), interceptorFactory); if (certKeyPair != null) { builder.useTls(certKeyPair.cert, certKeyPair.key); } @@ -301,11 +305,11 @@ public FlightServerTestExtension build() { * A middleware to handle with the cookies in the server. It is used to test if cookies are being * sent properly. */ - static class MiddlewareCookie implements FlightServerMiddleware { + static class InterceptorMiddleware implements FlightServerMiddleware { private final Factory factory; - public MiddlewareCookie(Factory factory) { + public InterceptorMiddleware(Factory factory) { this.factory = factory; } @@ -323,22 +327,33 @@ public void onCallCompleted(CallStatus callStatus) {} public void onCallErrored(Throwable throwable) {} /** A factory for the MiddlewareCookie. */ - static class Factory implements FlightServerMiddleware.Factory { + static class Factory implements FlightServerMiddleware.Factory { + private final Map receivedCallHeaders = new HashMap<>(); private boolean receivedCookieHeader = false; private String cookie; @Override - public MiddlewareCookie onCallStarted( + public InterceptorMiddleware onCallStarted( CallInfo callInfo, CallHeaders callHeaders, RequestContext requestContext) { cookie = callHeaders.get("Cookie"); receivedCookieHeader = null != cookie; - return new MiddlewareCookie(this); + + receivedCallHeaders.put(callInfo.method(), callHeaders); + return new InterceptorMiddleware(this); } public String getCookie() { return cookie; } + + public String getHeader(FlightMethod method, String key) { + CallHeaders headers = receivedCallHeaders.get(method); + if (headers == null) { + return null; + } + return headers.get(key); + } } } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerBuilderTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerBuilderTest.java index 6524eaf39a..a60a71f23d 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerBuilderTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerBuilderTest.java @@ -149,6 +149,7 @@ public void testDefaults() { assertEquals(Optional.empty(), builder.catalog); assertNull(builder.flightClientCache); assertNull(builder.connectTimeout); + assertNull(builder.driverVersion); } @Test From 16427331575f016b490002fc81dbe22a27d7852f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Thu, 31 Jul 2025 11:32:29 +0100 Subject: [PATCH 3/4] chore: remove VisibleForTesting --- .../arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 971a3f3f7c..0a35f8d8a9 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 @@ -549,7 +549,7 @@ public FlightInfo getCrossReference( /** Builder for {@link ArrowFlightSqlClientHandler}. */ public static final class Builder { - @VisibleForTesting static final String USER_AGENT_TEMPLATE = "JDBC Flight SQL Driver %s"; + static final String USER_AGENT_TEMPLATE = "JDBC Flight SQL Driver %s"; static final String DEFAULT_VERSION = "(unknown or development build)"; private final Set middlewareFactories = new HashSet<>(); From 82c55ae2c3e5436f3ec15d097ada242dc43a688a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Tue, 5 Aug 2025 14:58:47 +0100 Subject: [PATCH 4/4] remove useragent customization from NettyClientBuilder --- .../org/apache/arrow/flight/grpc/NettyClientBuilder.java | 9 +-------- .../driver/jdbc/client/ArrowFlightSqlClientHandler.java | 4 +++- 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/NettyClientBuilder.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/NettyClientBuilder.java index d237812051..42cdaac016 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/NettyClientBuilder.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/NettyClientBuilder.java @@ -60,7 +60,6 @@ public class NettyClientBuilder { protected String overrideHostname = null; protected List middleware = new ArrayList<>(); protected boolean verifyServer = true; - protected String userAgent; public NettyClientBuilder() {} @@ -131,11 +130,6 @@ public NettyClientBuilder verifyServer(boolean verifyServer) { return this; } - public NettyClientBuilder userAgent(String userAgent) { - this.userAgent = userAgent; - return this; - } - /** Create the client from this builder. */ public NettyChannelBuilder build() { final NettyChannelBuilder builder; @@ -232,8 +226,7 @@ public NettyChannelBuilder build() { builder .maxTraceEvents(MAX_CHANNEL_TRACE_EVENTS) .maxInboundMessageSize(maxInboundMessageSize) - .maxInboundMetadataSize(maxInboundMessageSize) - .userAgent(userAgent); + .maxInboundMetadataSize(maxInboundMessageSize); return builder; } } 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 0a35f8d8a9..a3f6900373 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 @@ -937,7 +937,6 @@ public ArrowFlightSqlClientHandler build() throws SQLException { if (driverVersion != null && driverVersion.versionString != null) { userAgent = String.format(USER_AGENT_TEMPLATE, driverVersion.versionString); } - clientBuilder.userAgent(userAgent); buildTimeMiddlewareFactories.add(new ClientCookieMiddleware.Factory()); buildTimeMiddlewareFactories.forEach(clientBuilder::intercept); @@ -973,6 +972,9 @@ public ArrowFlightSqlClientHandler build() throws SQLException { } NettyChannelBuilder channelBuilder = clientBuilder.build(); + + channelBuilder.userAgent(userAgent); + if (connectTimeout != null) { channelBuilder.withOption( ChannelOption.CONNECT_TIMEOUT_MILLIS, (int) connectTimeout.toMillis());