test: assert user-agent is correctly built

Author: xborder

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());
     }
   }
 }

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);
+    }
+  }
 }

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<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);
+      }
     }
   }
 }

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