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