This is an automated email from the ASF dual-hosted git repository.

lidavidm 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 a6245ab8 GH-804: Prepend JDBC FlightSQL version to user agent (#806)
a6245ab8 is described below

commit a6245ab87c782e2b873526acbd876cdf96d1f0ee
Author: Hélder Gregório <[email protected]>
AuthorDate: Thu Aug 7 01:29:41 2025 +0100

    GH-804: Prepend JDBC FlightSQL version to user agent (#806)
    
    ## What's Changed
    
    * Driver version is passed on to NettyBuilderClient to append it to the
    user-agent header
    * NettyBuilderClient now prepends `JDBC Flight SQL Client <version>` to
    the header (e.g. `JDBC Flight SQL Client 19.0.0-SNAPSHOT
    grpc-java-netty/1.73.0`)
    
    Closes #804.
---
 .../arrow/driver/jdbc/ArrowFlightConnection.java   |  9 ++++-
 .../jdbc/client/ArrowFlightSqlClientHandler.java   | 27 +++++++++++++
 .../jdbc/ArrowFlightJdbcConnectionCookieTest.java  |  4 +-
 .../apache/arrow/driver/jdbc/ConnectionTest.java   | 46 ++++++++++++++++++++++
 .../driver/jdbc/FlightServerTestExtension.java     | 33 +++++++++++-----
 .../ArrowFlightSqlClientHandlerBuilderTest.java    |  1 +
 6 files changed, 107 insertions(+), 13 deletions(-)

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 747287ed..f6f17770 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.AutoCloseables;
 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 @@ public final class ArrowFlightConnection extends 
AvaticaConnection {
       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 @@ public final class ArrowFlightConnection extends 
AvaticaConnection {
           .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 17c2c16e..a3f69003 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.Preconditions;
 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 final class ArrowFlightSqlClientHandler implements 
AutoCloseable {
 
   /** Builder for {@link ArrowFlightSqlClientHandler}. */
   public static final class Builder {
+    static final String USER_AGENT_TEMPLATE = "JDBC Flight SQL Driver %s";
+    static final String DEFAULT_VERSION = "(unknown or development build)";
+
     private final Set<FlightClientMiddleware.Factory> middlewareFactories = 
new HashSet<>();
     private final Set<CallOption> options = new HashSet<>();
     private String host;
@@ -597,6 +601,8 @@ public final class ArrowFlightSqlClientHandler implements 
AutoCloseable {
     @VisibleForTesting
     ClientCookieMiddleware.Factory cookieFactory = new 
ClientCookieMiddleware.Factory();
 
+    DriverVersion driverVersion;
+
     public Builder() {}
 
     /**
@@ -631,6 +637,8 @@ public final class ArrowFlightSqlClientHandler implements 
AutoCloseable {
       if (original.retainAuth) {
         this.authFactory = original.authFactory;
       }
+
+      this.driverVersion = original.driverVersion;
     }
 
     /**
@@ -879,6 +887,17 @@ public final class ArrowFlightSqlClientHandler implements 
AutoCloseable {
       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,11 @@ public final class ArrowFlightSqlClientHandler implements 
AutoCloseable {
         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);
+        }
+
         buildTimeMiddlewareFactories.add(new ClientCookieMiddleware.Factory());
         buildTimeMiddlewareFactories.forEach(clientBuilder::intercept);
         if (useEncryption) {
@@ -948,6 +972,9 @@ public final class ArrowFlightSqlClientHandler implements 
AutoCloseable {
         }
 
         NettyChannelBuilder channelBuilder = clientBuilder.build();
+
+        channelBuilder.userAgent(userAgent);
+
         if (connectTimeout != null) {
           channelBuilder.withOption(
               ChannelOption.CONNECT_TIMEOUT_MILLIS, (int) 
connectTimeout.toMillis());
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 1977b613..7127c7fc 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 class ArrowFlightJdbcConnectionCookieTest {
         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 8e872a11..72e4b222 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.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.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.util.AutoCloseables;
@@ -576,4 +577,49 @@ public class ConnectionTest {
       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 aa586651..db043805 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.Connection;
 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.driver.jdbc.utils.ArrowFlightConnectionConfigImpl;
 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 @@ public class FlightServerTestExtension
     properties.put("useEncryption", useEncryption);
   }
 
-  public MiddlewareCookie.Factory getMiddlewareCookieFactory() {
-    return middlewareCookieFactory;
+  public InterceptorMiddleware.Factory getInterceptorFactory() {
+    return interceptorFactory;
   }
 
   @FunctionalInterface
@@ -143,7 +147,7 @@ public class FlightServerTestExtension
     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 class FlightServerTestExtension
    * 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 class FlightServerTestExtension
     public void onCallErrored(Throwable throwable) {}
 
     /** A factory for the MiddlewareCookie. */
-    static class Factory implements 
FlightServerMiddleware.Factory<MiddlewareCookie> {
+    static class Factory implements 
FlightServerMiddleware.Factory<InterceptorMiddleware> {
 
+      private final Map<FlightMethod, CallHeaders> 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 6524eaf3..a60a71f2 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 class ArrowFlightSqlClientHandlerBuilderTest {
     assertEquals(Optional.empty(), builder.catalog);
     assertNull(builder.flightClientCache);
     assertNull(builder.connectTimeout);
+    assertNull(builder.driverVersion);
   }
 
   @Test

Reply via email to