From 1056424a3d88d061afa483c22bc504ea7528ae62 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andr=C3=A9=20EIckler?= <797483+eickler@users.noreply.github.com> Date: Sun, 4 Jan 2026 11:16:57 +0100 Subject: [PATCH 1/2] test: Test for Connection.setCatalog() --- .../arrow/driver/jdbc/ConnectionTest.java | 383 ++++++++++-------- .../jdbc/utils/MockFlightSqlProducer.java | 155 +++---- 2 files changed, 289 insertions(+), 249 deletions(-) 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 72e4b222a3..7ec40bd600 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.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; @@ -43,20 +47,20 @@ /** Tests for {@link Connection}. */ public class ConnectionTest { - @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION; + @RegisterExtension + public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION; private static final MockFlightSqlProducer PRODUCER = new MockFlightSqlProducer(); private static final String userTest = "user1"; private static final String passTest = "pass1"; static { - UserPasswordAuthentication authentication = - new UserPasswordAuthentication.Builder().user(userTest, passTest).build(); - - FLIGHT_SERVER_TEST_EXTENSION = - new FlightServerTestExtension.Builder() - .authentication(authentication) - .producer(PRODUCER) - .build(); + UserPasswordAuthentication authentication = new UserPasswordAuthentication.Builder().user(userTest, passTest) + .build(); + + FLIGHT_SERVER_TEST_EXTENSION = new FlightServerTestExtension.Builder() + .authentication(authentication) + .producer(PRODUCER) + .build(); } private BufferAllocator allocator; @@ -73,7 +77,8 @@ public void tearDown() throws Exception { } /** - * Checks if an unencrypted connection can be established successfully when the provided valid + * Checks if an unencrypted connection can be established successfully when the + * provided valid * credentials. * * @throws SQLException on error. @@ -90,19 +95,19 @@ public void testUnencryptedConnectionShouldOpenSuccessfullyWhenProvidedValidCred 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)) { + try (Connection connection = DriverManager.getConnection( + "jdbc:arrow-flight-sql://" + + FLIGHT_SERVER_TEST_EXTENSION.getHost() + + ":" + + FLIGHT_SERVER_TEST_EXTENSION.getPort(), + properties)) { assertTrue(connection.isValid(300)); } } /** - * Checks if a token is provided it takes precedence over username/pass. In this case, the + * Checks if a token is provided it takes precedence over username/pass. In this + * case, the * connection should fail if a token is passed in. */ @Test @@ -117,25 +122,24 @@ public void testTokenOverridesUsernameAndPasswordAuth() { properties.put(ArrowFlightConnectionProperty.TOKEN.camelName(), "token"); properties.put("useEncryption", false); - SQLException e = - assertThrows( - SQLException.class, - () -> { - try (Connection conn = - DriverManager.getConnection( - "jdbc:arrow-flight-sql://" - + FLIGHT_SERVER_TEST_EXTENSION.getHost() - + ":" - + FLIGHT_SERVER_TEST_EXTENSION.getPort(), - properties)) { - fail(); - } - }); + SQLException e = assertThrows( + SQLException.class, + () -> { + try (Connection conn = DriverManager.getConnection( + "jdbc:arrow-flight-sql://" + + FLIGHT_SERVER_TEST_EXTENSION.getHost() + + ":" + + FLIGHT_SERVER_TEST_EXTENSION.getPort(), + properties)) { + fail(); + } + }); assertTrue(e.getMessage().contains("UNAUTHENTICATED")); } /** - * Checks if the exception SQLException is thrown when trying to establish a connection without a + * Checks if the exception SQLException is thrown when trying to establish a + * connection without a * host. * * @throws SQLException on error. @@ -165,22 +169,22 @@ public void testUnencryptedConnectionWithEmptyHost() throws Exception { @Test public void testGetBasicClientAuthenticatedShouldOpenConnection() throws Exception { - try (ArrowFlightSqlClientHandler client = - new ArrowFlightSqlClientHandler.Builder() - .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) - .withPort(FLIGHT_SERVER_TEST_EXTENSION.getPort()) - .withEncryption(false) - .withUsername(userTest) - .withPassword(passTest) - .withBufferAllocator(allocator) - .build()) { + try (ArrowFlightSqlClientHandler client = new ArrowFlightSqlClientHandler.Builder() + .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) + .withPort(FLIGHT_SERVER_TEST_EXTENSION.getPort()) + .withEncryption(false) + .withUsername(userTest) + .withPassword(passTest) + .withBufferAllocator(allocator) + .build()) { assertNotNull(client); } } /** - * Checks if the exception IllegalArgumentException is thrown when trying to establish an + * Checks if the exception IllegalArgumentException is thrown when trying to + * establish an * unencrypted connection providing with an invalid port. * * @throws SQLException on error. @@ -193,8 +197,7 @@ public void testUnencryptedConnectionProvidingInvalidPort() throws Exception { properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), false); - final String invalidUrl = - "jdbc:arrow-flight-sql://" + FLIGHT_SERVER_TEST_EXTENSION.getHost() + ":" + 65537; + final String invalidUrl = "jdbc:arrow-flight-sql://" + FLIGHT_SERVER_TEST_EXTENSION.getHost() + ":" + 65537; assertThrows( SQLException.class, @@ -213,18 +216,18 @@ public void testUnencryptedConnectionProvidingInvalidPort() throws Exception { @Test public void testGetBasicClientNoAuthShouldOpenConnection() throws Exception { - try (ArrowFlightSqlClientHandler client = - new ArrowFlightSqlClientHandler.Builder() - .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) - .withBufferAllocator(allocator) - .withEncryption(false) - .build()) { + try (ArrowFlightSqlClientHandler client = new ArrowFlightSqlClientHandler.Builder() + .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) + .withBufferAllocator(allocator) + .withEncryption(false) + .build()) { assertNotNull(client); } } /** - * Checks if an unencrypted connection can be established successfully when not providing + * Checks if an unencrypted connection can be established successfully when not + * providing * credentials. * * @throws SQLException on error. @@ -237,14 +240,14 @@ public void testUnencryptedConnectionShouldOpenSuccessfullyWithoutAuthentication properties.put( ArrowFlightConnectionProperty.PORT.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getPort()); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), false); - try (Connection connection = - DriverManager.getConnection("jdbc:arrow-flight-sql://localhost:32010", properties)) { + try (Connection connection = DriverManager.getConnection("jdbc:arrow-flight-sql://localhost:32010", properties)) { assertTrue(connection.isValid(300)); } } /** - * Check if an unencrypted connection throws an exception when provided with invalid credentials. + * Check if an unencrypted connection throws an exception when provided with + * invalid credentials. * * @throws SQLException The exception expected to be thrown. */ @@ -264,15 +267,16 @@ public void testUnencryptedConnectionShouldThrowExceptionWhenProvidedWithInvalid assertThrows( SQLException.class, () -> { - try (Connection ignored = - DriverManager.getConnection("jdbc:arrow-flight-sql://localhost:32010", properties)) { + try ( + Connection ignored = DriverManager.getConnection("jdbc:arrow-flight-sql://localhost:32010", properties)) { fail(); } }); } /** - * Check if an non-encrypted connection can be established successfully when connecting through + * Check if an non-encrypted connection can be established successfully when + * connecting through * the DriverManager using just a connection url. * * @throws Exception on error. @@ -282,25 +286,25 @@ public void testTLSConnectionPropertyFalseCorrectCastUrlWithDriverManager() thro final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); - try (Connection connection = - DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&useEncryption=false", - FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest))) { + try (Connection connection = DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&useEncryption=false", + FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest))) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when connecting through - * the DriverManager using a connection url and properties with String K-V pairs. + * Check if an non-encrypted connection can be established successfully when + * connecting through + * the DriverManager using a connection url and properties with String K-V + * pairs. * * @throws Exception on error. */ @Test - public void - testTLSConnectionPropertyFalseCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() - throws Exception { + public void testTLSConnectionPropertyFalseCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() + throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); @@ -310,18 +314,19 @@ public void testTLSConnectionPropertyFalseCorrectCastUrlWithDriverManager() thro properties.setProperty(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.setProperty(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), "false"); - try (Connection connection = - DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), - properties)) { + try (Connection connection = DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), + properties)) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when connecting through - * the DriverManager using a connection url and properties with Object K-V pairs. + * Check if an non-encrypted connection can be established successfully when + * connecting through + * the DriverManager using a connection url and properties with Object K-V + * pairs. * * @throws Exception on error. */ @@ -336,18 +341,19 @@ public void testTLSConnectionPropertyFalseCorrectCastUrlAndPropertiesUsingPutWit properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), false); - try (Connection connection = - DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), - properties)) { + try (Connection connection = DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), + properties)) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when connecting through - * the DriverManager using just a connection url and using 0 and 1 as ssl values. + * Check if an non-encrypted connection can be established successfully when + * connecting through + * the DriverManager using just a connection url and using 0 and 1 as ssl + * values. * * @throws Exception on error. */ @@ -357,26 +363,26 @@ public void testTLSConnectionPropertyFalseIntegerCorrectCastUrlWithDriverManager final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); - try (Connection connection = - DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&useEncryption=0", - FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest))) { + try (Connection connection = DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&useEncryption=0", + FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest))) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when connecting through - * the DriverManager using a connection url and properties with String K-V pairs and using 0 and 1 + * Check if an non-encrypted connection can be established successfully when + * connecting through + * the DriverManager using a connection url and properties with String K-V pairs + * and using 0 and 1 * as ssl values. * * @throws Exception on error. */ @Test - public void - testTLSConnectionPropertyFalseIntegerCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() - throws Exception { + public void testTLSConnectionPropertyFalseIntegerCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() + throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); @@ -385,26 +391,26 @@ public void testTLSConnectionPropertyFalseIntegerCorrectCastUrlWithDriverManager properties.setProperty(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.setProperty(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), "0"); - try (Connection connection = - DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), - properties)) { + try (Connection connection = DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), + properties)) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when connecting through - * the DriverManager using a connection url and properties with Object K-V pairs and using 0 and 1 + * Check if an non-encrypted connection can be established successfully when + * connecting through + * the DriverManager using a connection url and properties with Object K-V pairs + * and using 0 and 1 * as ssl values. * * @throws Exception on error. */ @Test - public void - testTLSConnectionPropertyFalseIntegerCorrectCastUrlAndPropertiesUsingPutWithDriverManager() - throws Exception { + public void testTLSConnectionPropertyFalseIntegerCorrectCastUrlAndPropertiesUsingPutWithDriverManager() + throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); @@ -413,17 +419,17 @@ public void testTLSConnectionPropertyFalseIntegerCorrectCastUrlWithDriverManager properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), 0); - try (Connection connection = - DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), - properties)) { + try (Connection connection = DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), + properties)) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when connecting through + * Check if an non-encrypted connection can be established successfully when + * connecting through * the DriverManager using just a connection url. * * @throws Exception on error. @@ -434,26 +440,26 @@ public void testThreadPoolSizeConnectionPropertyCorrectCastUrlWithDriverManager( final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); - try (Connection connection = - DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&threadPoolSize=1&useEncryption=%s", - FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest, false))) { + try (Connection connection = DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&threadPoolSize=1&useEncryption=%s", + FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest, false))) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when connecting through - * the DriverManager using a connection url and properties with String K-V pairs and using 0 and 1 + * Check if an non-encrypted connection can be established successfully when + * connecting through + * the DriverManager using a connection url and properties with String K-V pairs + * and using 0 and 1 * as ssl values. * * @throws Exception on error. */ @Test - public void - testThreadPoolSizeConnectionPropertyCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() - throws Exception { + public void testThreadPoolSizeConnectionPropertyCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() + throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); @@ -463,26 +469,26 @@ public void testThreadPoolSizeConnectionPropertyCorrectCastUrlWithDriverManager( properties.setProperty(ArrowFlightConnectionProperty.THREAD_POOL_SIZE.camelName(), "1"); properties.put("useEncryption", false); - try (Connection connection = - DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), - properties)) { + try (Connection connection = DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), + properties)) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when connecting through - * the DriverManager using a connection url and properties with Object K-V pairs and using 0 and 1 + * Check if an non-encrypted connection can be established successfully when + * connecting through + * the DriverManager using a connection url and properties with Object K-V pairs + * and using 0 and 1 * as ssl values. * * @throws Exception on error. */ @Test - public void - testThreadPoolSizeConnectionPropertyCorrectCastUrlAndPropertiesUsingPutWithDriverManager() - throws Exception { + public void testThreadPoolSizeConnectionPropertyCorrectCastUrlAndPropertiesUsingPutWithDriverManager() + throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); @@ -492,17 +498,17 @@ public void testThreadPoolSizeConnectionPropertyCorrectCastUrlWithDriverManager( properties.put(ArrowFlightConnectionProperty.THREAD_POOL_SIZE.camelName(), 1); properties.put("useEncryption", false); - try (Connection connection = - DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), - properties)) { + try (Connection connection = DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), + properties)) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when connecting through + * Check if an non-encrypted connection can be established successfully when + * connecting through * the DriverManager using just a connection url. * * @throws Exception on error. @@ -513,26 +519,26 @@ public void testPasswordConnectionPropertyIntegerCorrectCastUrlWithDriverManager final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); - try (Connection connection = - DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&useEncryption=%s", - FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest, false))) { + try (Connection connection = DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&useEncryption=%s", + FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest, false))) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when connecting through - * the DriverManager using a connection url and properties with String K-V pairs and using 0 and 1 + * Check if an non-encrypted connection can be established successfully when + * connecting through + * the DriverManager using a connection url and properties with String K-V pairs + * and using 0 and 1 * as ssl values. * * @throws Exception on error. */ @Test - public void - testPasswordConnectionPropertyIntegerCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() - throws Exception { + public void testPasswordConnectionPropertyIntegerCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() + throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); @@ -541,26 +547,26 @@ public void testPasswordConnectionPropertyIntegerCorrectCastUrlWithDriverManager properties.setProperty(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put("useEncryption", false); - try (Connection connection = - DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), - properties)) { + try (Connection connection = DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), + properties)) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when connecting through - * the DriverManager using a connection url and properties with Object K-V pairs and using 0 and 1 + * Check if an non-encrypted connection can be established successfully when + * connecting through + * the DriverManager using a connection url and properties with Object K-V pairs + * and using 0 and 1 * as ssl values. * * @throws Exception on error. */ @Test - public void - testPasswordConnectionPropertyIntegerCorrectCastUrlAndPropertiesUsingPutWithDriverManager() - throws Exception { + public void testPasswordConnectionPropertyIntegerCorrectCastUrlAndPropertiesUsingPutWithDriverManager() + throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); @@ -569,17 +575,17 @@ public void testPasswordConnectionPropertyIntegerCorrectCastUrlWithDriverManager properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put("useEncryption", false); - try (Connection connection = - DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), - properties)) { + try (Connection connection = DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), + properties)) { assertTrue(connection.isValid(0)); } } /** - * Test that the JDBC driver properly integrates driver version into client handler. + * Test that the JDBC driver properly integrates driver version into client + * handler. * * @throws Exception on error. */ @@ -597,24 +603,22 @@ public void testJdbcDriverVersionIntegration() throws Exception { // 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)) { + 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 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 + 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( @@ -622,4 +626,35 @@ public void testJdbcDriverVersionIntegration() throws Exception { "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 options = PRODUCER.getSessionOptions(); + assertTrue(options.containsKey("catalog")); + String actualCatalog = options + .get("catalog") + .acceptVisitor( + new NoOpSessionOptionValueVisitor() { + @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 a8874c4869..12f62f6729 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.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; @@ -91,10 +94,8 @@ public final class MockFlightSqlProducer implements FlightSqlProducer { private final Map>> queryResults = new HashMap<>(); private final Map> selectResultProviders = new HashMap<>(); private final Map preparedStatements = new HashMap<>(); - private final Map> catalogQueriesResults = - new HashMap<>(); - private final Map>> - updateResultProviders = new HashMap<>(); + private final Map> catalogQueriesResults = new HashMap<>(); + private final Map>> updateResultProviders = new HashMap<>(); private final SqlInfoBuilder sqlInfoBuilder = new SqlInfoBuilder(); private final Map parameterSchemas = new HashMap<>(); private final Map>> expectedParameterValues = new HashMap<>(); @@ -130,8 +131,8 @@ public static ByteBuffer serializeSchema(final Schema schema) { /** * Registers a new {@link StatementType#SELECT} SQL query. * - * @param sqlCommand the SQL command under which to register the new query. - * @param schema the schema to use for the query result. + * @param sqlCommand the SQL command under which to register the new query. + * @param schema the schema to use for the query result. * @param resultProviders the result provider for this query. */ public void addSelectQuery( @@ -139,10 +140,9 @@ public void addSelectQuery( final Schema schema, final List> resultProviders) { final int providers = resultProviders.size(); - final List uuids = - IntStream.range(0, providers) - .mapToObj(index -> new UUID(sqlCommand.hashCode(), Integer.hashCode(index))) - .collect(toList()); + final List uuids = IntStream.range(0, providers) + .mapToObj(index -> new UUID(sqlCommand.hashCode(), Integer.hashCode(index))) + .collect(toList()); queryResults.put(sqlCommand, new SimpleImmutableEntry<>(schema, uuids)); IntStream.range(0, providers) .forEach( @@ -152,15 +152,14 @@ public void addSelectQuery( /** * Registers a new {@link StatementType#UPDATE} SQL query. * - * @param sqlCommand the SQL command. + * @param sqlCommand the SQL command. * @param updatedRows the number of rows affected. */ public void addUpdateQuery(final String sqlCommand, final long updatedRows) { addUpdateQuery( sqlCommand, (flightStream, putResultStreamListener) -> { - final DoPutUpdateResult result = - DoPutUpdateResult.newBuilder().setRecordCount(updatedRows).build(); + final DoPutUpdateResult result = DoPutUpdateResult.newBuilder().setRecordCount(updatedRows).build(); try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); final ArrowBuf buffer = allocator.buffer(result.getSerializedSize())) { buffer.writeBytes(result.toByteArray()); @@ -176,7 +175,8 @@ public void addUpdateQuery(final String sqlCommand, final long updatedRows) { /** * Adds a catalog query to the results. * - * @param message the {@link Message} corresponding to the catalog query request type to register. + * @param message the {@link Message} corresponding to the catalog query + * request type to register. * @param resultsProvider the results provider. */ public void addCatalogQuery( @@ -187,7 +187,7 @@ public void addCatalogQuery( /** * Registers a new {@link StatementType#UPDATE} SQL query. * - * @param sqlCommand the SQL command. + * @param sqlCommand the SQL command. * @param resultsProvider consumer for producing update results. */ void addUpdateQuery( @@ -211,13 +211,11 @@ public void createPreparedStatement( final CallContext callContext, final StreamListener listener) { try { - final ByteString preparedStatementHandle = - copyFrom(randomUUID().toString().getBytes(StandardCharsets.UTF_8)); + final ByteString preparedStatementHandle = copyFrom(randomUUID().toString().getBytes(StandardCharsets.UTF_8)); final String query = request.getQuery(); - final ActionCreatePreparedStatementResult.Builder resultBuilder = - ActionCreatePreparedStatementResult.newBuilder() - .setPreparedStatementHandle(preparedStatementHandle); + final ActionCreatePreparedStatementResult.Builder resultBuilder = ActionCreatePreparedStatementResult.newBuilder() + .setPreparedStatementHandle(preparedStatementHandle); final Entry> entry = queryResults.get(query); if (entry != null) { @@ -267,15 +265,13 @@ public FlightInfo getFlightInfoStatement( final CallContext callContext, final FlightDescriptor flightDescriptor) { final String query = commandStatementQuery.getQuery(); - final Entry> queryInfo = - Preconditions.checkNotNull( - queryResults.get(query), format("Query not registered: <%s>.", query)); - final List endpoints = - queryInfo.getValue().stream() - .map(TicketConversionUtils::getTicketBytesFromUuid) - .map(TicketConversionUtils::getTicketStatementQueryFromHandle) - .map(TicketConversionUtils::getEndpointFromMessage) - .collect(toList()); + final Entry> queryInfo = Preconditions.checkNotNull( + queryResults.get(query), format("Query not registered: <%s>.", query)); + final List endpoints = queryInfo.getValue().stream() + .map(TicketConversionUtils::getTicketBytesFromUuid) + .map(TicketConversionUtils::getTicketStatementQueryFromHandle) + .map(TicketConversionUtils::getEndpointFromMessage) + .collect(toList()); return FlightInfo.builder(queryInfo.getKey(), flightDescriptor, endpoints) .setAppMetadata("foo".getBytes(StandardCharsets.UTF_8)) .build(); @@ -286,22 +282,18 @@ public FlightInfo getFlightInfoPreparedStatement( final CommandPreparedStatementQuery commandPreparedStatementQuery, final CallContext callContext, final FlightDescriptor flightDescriptor) { - final ByteString preparedStatementHandle = - commandPreparedStatementQuery.getPreparedStatementHandle(); - - final String query = - Preconditions.checkNotNull( - preparedStatements.get(preparedStatementHandle), - format("No query registered under handle: <%s>.", preparedStatementHandle)); - final Entry> queryInfo = - Preconditions.checkNotNull( - queryResults.get(query), format("Query not registered: <%s>.", query)); - final List endpoints = - queryInfo.getValue().stream() - .map(TicketConversionUtils::getTicketBytesFromUuid) - .map(TicketConversionUtils::getCommandPreparedStatementQueryFromHandle) - .map(TicketConversionUtils::getEndpointFromMessage) - .collect(toList()); + final ByteString preparedStatementHandle = commandPreparedStatementQuery.getPreparedStatementHandle(); + + final String query = Preconditions.checkNotNull( + preparedStatements.get(preparedStatementHandle), + format("No query registered under handle: <%s>.", preparedStatementHandle)); + final Entry> queryInfo = Preconditions.checkNotNull( + queryResults.get(query), format("Query not registered: <%s>.", query)); + final List endpoints = queryInfo.getValue().stream() + .map(TicketConversionUtils::getTicketBytesFromUuid) + .map(TicketConversionUtils::getCommandPreparedStatementQueryFromHandle) + .map(TicketConversionUtils::getEndpointFromMessage) + .collect(toList()); return FlightInfo.builder(queryInfo.getKey(), flightDescriptor, endpoints) .setAppMetadata("foo".getBytes(StandardCharsets.UTF_8)) .build(); @@ -313,9 +305,8 @@ public SchemaResult getSchemaStatement( final CallContext callContext, final FlightDescriptor flightDescriptor) { final String query = commandStatementQuery.getQuery(); - final Entry> queryInfo = - Preconditions.checkNotNull( - queryResults.get(query), format("Query not registered: <%s>.", query)); + final Entry> queryInfo = Preconditions.checkNotNull( + queryResults.get(query), format("Query not registered: <%s>.", query)); return new SchemaResult(queryInfo.getKey()); } @@ -327,9 +318,9 @@ public void getStreamStatement( final ServerStreamListener serverStreamListener) { final UUID uuid = UUID.fromString(ticketStatementQuery.getStatementHandle().toStringUtf8()); Preconditions.checkNotNull( - selectResultProviders.get(uuid), - "No consumer was registered for the specified UUID: <%s>.", - uuid) + selectResultProviders.get(uuid), + "No consumer was registered for the specified UUID: <%s>.", + uuid) .accept(serverStreamListener); } @@ -338,12 +329,11 @@ public void getStreamPreparedStatement( final CommandPreparedStatementQuery commandPreparedStatementQuery, final CallContext callContext, final ServerStreamListener serverStreamListener) { - final UUID uuid = - UUID.fromString(commandPreparedStatementQuery.getPreparedStatementHandle().toStringUtf8()); + final UUID uuid = UUID.fromString(commandPreparedStatementQuery.getPreparedStatementHandle().toStringUtf8()); Preconditions.checkNotNull( - selectResultProviders.get(uuid), - "No consumer was registered for the specified UUID: <%s>.", - uuid) + selectResultProviders.get(uuid), + "No consumer was registered for the specified UUID: <%s>.", + uuid) .accept(serverStreamListener); } @@ -355,10 +345,9 @@ public Runnable acceptPutStatement( final StreamListener streamListener) { return () -> { final String query = commandStatementUpdate.getQuery(); - final BiConsumer> resultProvider = - Preconditions.checkNotNull( - updateResultProviders.get(query), - format("No consumer found for query: <%s>.", query)); + final BiConsumer> resultProvider = Preconditions.checkNotNull( + updateResultProviders.get(query), + format("No consumer found for query: <%s>.", query)); resultProvider.accept(flightStream, streamListener); }; } @@ -392,8 +381,7 @@ private boolean validateParameters( Object actual = root.getVector(paramIndex).getObject(i); boolean matches; if (expected.getClass().isArray()) { - matches = - Arrays.equals((Object[]) expected, ((JsonStringArrayList) actual).toArray()); + matches = Arrays.equals((Object[]) expected, ((JsonStringArrayList) actual).toArray()); } else { matches = Objects.equals(expected, actual); } @@ -427,13 +415,13 @@ public Runnable acceptPutPreparedStatementUpdate( final FlightStream flightStream, final StreamListener streamListener) { final ByteString handle = commandPreparedStatementUpdate.getPreparedStatementHandle(); - final String query = - Preconditions.checkNotNull( - preparedStatements.get(handle), - format("No query registered under handle: <%s>.", handle)); + final String query = Preconditions.checkNotNull( + preparedStatements.get(handle), + format("No query registered under handle: <%s>.", handle)); if (validateParameters(query, flightStream, streamListener)) { - return () -> {}; + return () -> { + }; } return acceptPutStatement( @@ -450,13 +438,13 @@ public Runnable acceptPutPreparedStatementQuery( final FlightStream flightStream, final StreamListener streamListener) { final ByteString handle = commandPreparedStatementQuery.getPreparedStatementHandle(); - final String query = - Preconditions.checkNotNull( - preparedStatements.get(handle), - format("No query registered under handle: <%s>.", handle)); + final String query = Preconditions.checkNotNull( + preparedStatements.get(handle), + format("No query registered under handle: <%s>.", handle)); if (validateParameters(query, flightStream, streamListener)) { - return () -> {}; + return () -> { + }; } return streamListener::onCompleted; @@ -641,7 +629,8 @@ public void doAction(CallContext context, Action action, StreamListener } /** - * Clear the `actionTypeCounter` map and restore to its default state. Intended to be used in + * Clear the `actionTypeCounter` map and restore to its default state. Intended + * to be used in * tests. */ public void clearActionTypeCounter() { @@ -655,8 +644,8 @@ public Map getActionTypeCounter() { private void getStreamCatalogFunctions( final Message ticket, final ServerStreamListener serverStreamListener) { Preconditions.checkNotNull( - catalogQueriesResults.get(ticket), - format("Query not registered for ticket: <%s>", ticket)) + catalogQueriesResults.get(ticket), + format("Query not registered for ticket: <%s>", ticket)) .accept(serverStreamListener); } @@ -664,6 +653,22 @@ public SqlInfoBuilder getSqlInfoBuilder() { return sqlInfoBuilder; } + private final Map sessionOptions = new HashMap<>(); + + @Override + public void setSessionOptions( + final SetSessionOptionsRequest request, + final CallContext context, + final StreamListener listener) { + sessionOptions.putAll(request.getSessionOptions()); + listener.onNext(new SetSessionOptionsResult(Collections.emptyMap())); + listener.onCompleted(); + } + + public Map getSessionOptions() { + return sessionOptions; + } + private static final class TicketConversionUtils { private TicketConversionUtils() { // Prevent instantiation. From c09bc391eb9192c8db0133a0dad8054826d3bb33 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andr=C3=A9=20EIckler?= <797483+eickler@users.noreply.github.com> Date: Sun, 4 Jan 2026 11:28:06 +0100 Subject: [PATCH 2/2] fix: Connection.setCatalog() fixed --- .../driver/jdbc/ArrowFlightMetaImpl.java | 23 +- .../client/ArrowFlightSqlClientHandler.java | 80 ++-- .../arrow/driver/jdbc/ConnectionTest.java | 377 +++++++++--------- .../jdbc/utils/MockFlightSqlProducer.java | 136 ++++--- 4 files changed, 338 insertions(+), 278 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 21cc3e431f..64529b50c8 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 @@ static Signature newSignature(final String sql, Schema resultSetSchema, Schema p 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 ExecuteResult prepareAndExecute( 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 boolean syncResults( 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 @@ PreparedStatement getPreparedStatement(StatementHandle statementHandle) { 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 5dc7e0e2e9..666996cd95 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.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 List getStreams(final FlightInfo flightInfo) 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 exceptions = new ArrayList<>(); CloseableEndpointStreamPair stream = null; @@ -337,7 +342,8 @@ private boolean isBenignCloseException(FlightRuntimeException fre) { */ 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 interface PreparedStatement extends AutoCloseable { /** A connection is created with catalog set as a session option. */ private void setSetCatalogInSessionIfPresent() { if (catalog.isPresent()) { - final SetSessionOptionsRequest setSessionOptionRequest = - new SetSessionOptionsRequest( - ImmutableMap.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 errors = result.getErrors(); - for (Map.Entry error : errors.entrySet()) { + final Map errors = result.getErrors(); + for (final Map.Entry 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 static final class Builder { @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 Location getLocation() { * @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 buildTimeMiddlewareFactories = new HashSet<>(this.middlewareFactories); @@ -988,7 +1011,8 @@ public ArrowFlightSqlClientHandler build() throws SQLException { 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 ArrowFlightSqlClientHandler build() throws SQLException { allocator, channelBuilder.build(), clientBuilder.middleware()); final ArrayList 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 7ec40bd600..46762f3319 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 @@ -47,20 +47,20 @@ /** Tests for {@link Connection}. */ public class ConnectionTest { - @RegisterExtension - public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION; + @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION; private static final MockFlightSqlProducer PRODUCER = new MockFlightSqlProducer(); private static final String userTest = "user1"; private static final String passTest = "pass1"; static { - UserPasswordAuthentication authentication = new UserPasswordAuthentication.Builder().user(userTest, passTest) - .build(); - - FLIGHT_SERVER_TEST_EXTENSION = new FlightServerTestExtension.Builder() - .authentication(authentication) - .producer(PRODUCER) - .build(); + UserPasswordAuthentication authentication = + new UserPasswordAuthentication.Builder().user(userTest, passTest).build(); + + FLIGHT_SERVER_TEST_EXTENSION = + new FlightServerTestExtension.Builder() + .authentication(authentication) + .producer(PRODUCER) + .build(); } private BufferAllocator allocator; @@ -77,8 +77,7 @@ public void tearDown() throws Exception { } /** - * Checks if an unencrypted connection can be established successfully when the - * provided valid + * Checks if an unencrypted connection can be established successfully when the provided valid * credentials. * * @throws SQLException on error. @@ -95,19 +94,19 @@ public void testUnencryptedConnectionShouldOpenSuccessfullyWhenProvidedValidCred 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)) { + try (Connection connection = + DriverManager.getConnection( + "jdbc:arrow-flight-sql://" + + FLIGHT_SERVER_TEST_EXTENSION.getHost() + + ":" + + FLIGHT_SERVER_TEST_EXTENSION.getPort(), + properties)) { assertTrue(connection.isValid(300)); } } /** - * Checks if a token is provided it takes precedence over username/pass. In this - * case, the + * Checks if a token is provided it takes precedence over username/pass. In this case, the * connection should fail if a token is passed in. */ @Test @@ -122,24 +121,25 @@ public void testTokenOverridesUsernameAndPasswordAuth() { properties.put(ArrowFlightConnectionProperty.TOKEN.camelName(), "token"); properties.put("useEncryption", false); - SQLException e = assertThrows( - SQLException.class, - () -> { - try (Connection conn = DriverManager.getConnection( - "jdbc:arrow-flight-sql://" - + FLIGHT_SERVER_TEST_EXTENSION.getHost() - + ":" - + FLIGHT_SERVER_TEST_EXTENSION.getPort(), - properties)) { - fail(); - } - }); + SQLException e = + assertThrows( + SQLException.class, + () -> { + try (Connection conn = + DriverManager.getConnection( + "jdbc:arrow-flight-sql://" + + FLIGHT_SERVER_TEST_EXTENSION.getHost() + + ":" + + FLIGHT_SERVER_TEST_EXTENSION.getPort(), + properties)) { + fail(); + } + }); assertTrue(e.getMessage().contains("UNAUTHENTICATED")); } /** - * Checks if the exception SQLException is thrown when trying to establish a - * connection without a + * Checks if the exception SQLException is thrown when trying to establish a connection without a * host. * * @throws SQLException on error. @@ -169,22 +169,22 @@ public void testUnencryptedConnectionWithEmptyHost() throws Exception { @Test public void testGetBasicClientAuthenticatedShouldOpenConnection() throws Exception { - try (ArrowFlightSqlClientHandler client = new ArrowFlightSqlClientHandler.Builder() - .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) - .withPort(FLIGHT_SERVER_TEST_EXTENSION.getPort()) - .withEncryption(false) - .withUsername(userTest) - .withPassword(passTest) - .withBufferAllocator(allocator) - .build()) { + try (ArrowFlightSqlClientHandler client = + new ArrowFlightSqlClientHandler.Builder() + .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) + .withPort(FLIGHT_SERVER_TEST_EXTENSION.getPort()) + .withEncryption(false) + .withUsername(userTest) + .withPassword(passTest) + .withBufferAllocator(allocator) + .build()) { assertNotNull(client); } } /** - * Checks if the exception IllegalArgumentException is thrown when trying to - * establish an + * Checks if the exception IllegalArgumentException is thrown when trying to establish an * unencrypted connection providing with an invalid port. * * @throws SQLException on error. @@ -197,7 +197,8 @@ public void testUnencryptedConnectionProvidingInvalidPort() throws Exception { properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), false); - final String invalidUrl = "jdbc:arrow-flight-sql://" + FLIGHT_SERVER_TEST_EXTENSION.getHost() + ":" + 65537; + final String invalidUrl = + "jdbc:arrow-flight-sql://" + FLIGHT_SERVER_TEST_EXTENSION.getHost() + ":" + 65537; assertThrows( SQLException.class, @@ -216,18 +217,18 @@ public void testUnencryptedConnectionProvidingInvalidPort() throws Exception { @Test public void testGetBasicClientNoAuthShouldOpenConnection() throws Exception { - try (ArrowFlightSqlClientHandler client = new ArrowFlightSqlClientHandler.Builder() - .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) - .withBufferAllocator(allocator) - .withEncryption(false) - .build()) { + try (ArrowFlightSqlClientHandler client = + new ArrowFlightSqlClientHandler.Builder() + .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) + .withBufferAllocator(allocator) + .withEncryption(false) + .build()) { assertNotNull(client); } } /** - * Checks if an unencrypted connection can be established successfully when not - * providing + * Checks if an unencrypted connection can be established successfully when not providing * credentials. * * @throws SQLException on error. @@ -240,14 +241,14 @@ public void testUnencryptedConnectionShouldOpenSuccessfullyWithoutAuthentication properties.put( ArrowFlightConnectionProperty.PORT.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getPort()); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), false); - try (Connection connection = DriverManager.getConnection("jdbc:arrow-flight-sql://localhost:32010", properties)) { + try (Connection connection = + DriverManager.getConnection("jdbc:arrow-flight-sql://localhost:32010", properties)) { assertTrue(connection.isValid(300)); } } /** - * Check if an unencrypted connection throws an exception when provided with - * invalid credentials. + * Check if an unencrypted connection throws an exception when provided with invalid credentials. * * @throws SQLException The exception expected to be thrown. */ @@ -267,16 +268,15 @@ public void testUnencryptedConnectionShouldThrowExceptionWhenProvidedWithInvalid assertThrows( SQLException.class, () -> { - try ( - Connection ignored = DriverManager.getConnection("jdbc:arrow-flight-sql://localhost:32010", properties)) { + try (Connection ignored = + DriverManager.getConnection("jdbc:arrow-flight-sql://localhost:32010", properties)) { fail(); } }); } /** - * Check if an non-encrypted connection can be established successfully when - * connecting through + * Check if an non-encrypted connection can be established successfully when connecting through * the DriverManager using just a connection url. * * @throws Exception on error. @@ -286,25 +286,25 @@ public void testTLSConnectionPropertyFalseCorrectCastUrlWithDriverManager() thro final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); - try (Connection connection = DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&useEncryption=false", - FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest))) { + try (Connection connection = + DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&useEncryption=false", + FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest))) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when - * connecting through - * the DriverManager using a connection url and properties with String K-V - * pairs. + * Check if an non-encrypted connection can be established successfully when connecting through + * the DriverManager using a connection url and properties with String K-V pairs. * * @throws Exception on error. */ @Test - public void testTLSConnectionPropertyFalseCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() - throws Exception { + public void + testTLSConnectionPropertyFalseCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() + throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); @@ -314,19 +314,18 @@ public void testTLSConnectionPropertyFalseCorrectCastUrlAndPropertiesUsingSetPro properties.setProperty(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.setProperty(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), "false"); - try (Connection connection = DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), - properties)) { + try (Connection connection = + DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), + properties)) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when - * connecting through - * the DriverManager using a connection url and properties with Object K-V - * pairs. + * Check if an non-encrypted connection can be established successfully when connecting through + * the DriverManager using a connection url and properties with Object K-V pairs. * * @throws Exception on error. */ @@ -341,19 +340,18 @@ public void testTLSConnectionPropertyFalseCorrectCastUrlAndPropertiesUsingPutWit properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), false); - try (Connection connection = DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), - properties)) { + try (Connection connection = + DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), + properties)) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when - * connecting through - * the DriverManager using just a connection url and using 0 and 1 as ssl - * values. + * Check if an non-encrypted connection can be established successfully when connecting through + * the DriverManager using just a connection url and using 0 and 1 as ssl values. * * @throws Exception on error. */ @@ -363,26 +361,26 @@ public void testTLSConnectionPropertyFalseIntegerCorrectCastUrlWithDriverManager final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); - try (Connection connection = DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&useEncryption=0", - FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest))) { + try (Connection connection = + DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&useEncryption=0", + FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest))) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when - * connecting through - * the DriverManager using a connection url and properties with String K-V pairs - * and using 0 and 1 + * Check if an non-encrypted connection can be established successfully when connecting through + * the DriverManager using a connection url and properties with String K-V pairs and using 0 and 1 * as ssl values. * * @throws Exception on error. */ @Test - public void testTLSConnectionPropertyFalseIntegerCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() - throws Exception { + public void + testTLSConnectionPropertyFalseIntegerCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() + throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); @@ -391,26 +389,26 @@ public void testTLSConnectionPropertyFalseIntegerCorrectCastUrlAndPropertiesUsin properties.setProperty(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.setProperty(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), "0"); - try (Connection connection = DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), - properties)) { + try (Connection connection = + DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), + properties)) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when - * connecting through - * the DriverManager using a connection url and properties with Object K-V pairs - * and using 0 and 1 + * Check if an non-encrypted connection can be established successfully when connecting through + * the DriverManager using a connection url and properties with Object K-V pairs and using 0 and 1 * as ssl values. * * @throws Exception on error. */ @Test - public void testTLSConnectionPropertyFalseIntegerCorrectCastUrlAndPropertiesUsingPutWithDriverManager() - throws Exception { + public void + testTLSConnectionPropertyFalseIntegerCorrectCastUrlAndPropertiesUsingPutWithDriverManager() + throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); @@ -419,17 +417,17 @@ public void testTLSConnectionPropertyFalseIntegerCorrectCastUrlAndPropertiesUsin properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), 0); - try (Connection connection = DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), - properties)) { + try (Connection connection = + DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), + properties)) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when - * connecting through + * Check if an non-encrypted connection can be established successfully when connecting through * the DriverManager using just a connection url. * * @throws Exception on error. @@ -440,26 +438,26 @@ public void testThreadPoolSizeConnectionPropertyCorrectCastUrlWithDriverManager( final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); - try (Connection connection = DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&threadPoolSize=1&useEncryption=%s", - FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest, false))) { + try (Connection connection = + DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&threadPoolSize=1&useEncryption=%s", + FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest, false))) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when - * connecting through - * the DriverManager using a connection url and properties with String K-V pairs - * and using 0 and 1 + * Check if an non-encrypted connection can be established successfully when connecting through + * the DriverManager using a connection url and properties with String K-V pairs and using 0 and 1 * as ssl values. * * @throws Exception on error. */ @Test - public void testThreadPoolSizeConnectionPropertyCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() - throws Exception { + public void + testThreadPoolSizeConnectionPropertyCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() + throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); @@ -469,26 +467,26 @@ public void testThreadPoolSizeConnectionPropertyCorrectCastUrlAndPropertiesUsing properties.setProperty(ArrowFlightConnectionProperty.THREAD_POOL_SIZE.camelName(), "1"); properties.put("useEncryption", false); - try (Connection connection = DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), - properties)) { + try (Connection connection = + DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), + properties)) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when - * connecting through - * the DriverManager using a connection url and properties with Object K-V pairs - * and using 0 and 1 + * Check if an non-encrypted connection can be established successfully when connecting through + * the DriverManager using a connection url and properties with Object K-V pairs and using 0 and 1 * as ssl values. * * @throws Exception on error. */ @Test - public void testThreadPoolSizeConnectionPropertyCorrectCastUrlAndPropertiesUsingPutWithDriverManager() - throws Exception { + public void + testThreadPoolSizeConnectionPropertyCorrectCastUrlAndPropertiesUsingPutWithDriverManager() + throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); @@ -498,17 +496,17 @@ public void testThreadPoolSizeConnectionPropertyCorrectCastUrlAndPropertiesUsing properties.put(ArrowFlightConnectionProperty.THREAD_POOL_SIZE.camelName(), 1); properties.put("useEncryption", false); - try (Connection connection = DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), - properties)) { + try (Connection connection = + DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), + properties)) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when - * connecting through + * Check if an non-encrypted connection can be established successfully when connecting through * the DriverManager using just a connection url. * * @throws Exception on error. @@ -519,26 +517,26 @@ public void testPasswordConnectionPropertyIntegerCorrectCastUrlWithDriverManager final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); - try (Connection connection = DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&useEncryption=%s", - FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest, false))) { + try (Connection connection = + DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&useEncryption=%s", + FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest, false))) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when - * connecting through - * the DriverManager using a connection url and properties with String K-V pairs - * and using 0 and 1 + * Check if an non-encrypted connection can be established successfully when connecting through + * the DriverManager using a connection url and properties with String K-V pairs and using 0 and 1 * as ssl values. * * @throws Exception on error. */ @Test - public void testPasswordConnectionPropertyIntegerCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() - throws Exception { + public void + testPasswordConnectionPropertyIntegerCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() + throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); @@ -547,26 +545,26 @@ public void testPasswordConnectionPropertyIntegerCorrectCastUrlAndPropertiesUsin properties.setProperty(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put("useEncryption", false); - try (Connection connection = DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), - properties)) { + try (Connection connection = + DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), + properties)) { assertTrue(connection.isValid(0)); } } /** - * Check if an non-encrypted connection can be established successfully when - * connecting through - * the DriverManager using a connection url and properties with Object K-V pairs - * and using 0 and 1 + * Check if an non-encrypted connection can be established successfully when connecting through + * the DriverManager using a connection url and properties with Object K-V pairs and using 0 and 1 * as ssl values. * * @throws Exception on error. */ @Test - public void testPasswordConnectionPropertyIntegerCorrectCastUrlAndPropertiesUsingPutWithDriverManager() - throws Exception { + public void + testPasswordConnectionPropertyIntegerCorrectCastUrlAndPropertiesUsingPutWithDriverManager() + throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); @@ -575,17 +573,17 @@ public void testPasswordConnectionPropertyIntegerCorrectCastUrlAndPropertiesUsin properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put("useEncryption", false); - try (Connection connection = DriverManager.getConnection( - String.format( - "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), - properties)) { + try (Connection connection = + DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), + properties)) { assertTrue(connection.isValid(0)); } } /** - * Test that the JDBC driver properly integrates driver version into client - * handler. + * Test that the JDBC driver properly integrates driver version into client handler. * * @throws Exception on error. */ @@ -603,20 +601,23 @@ public void testJdbcDriverVersionIntegration() throws Exception { // 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)) { + 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 actualUserAgent = + FLIGHT_SERVER_TEST_EXTENSION + .getInterceptorFactory() + .getHeader(FlightMethod.HANDSHAKE, "user-agent"); - var expectedUserAgent = "JDBC Flight SQL Driver " + driverVersion.getDriverVersion().versionString; + var expectedUserAgent = + "JDBC Flight SQL Driver " + driverVersion.getDriverVersion().versionString; // Driver appends version to grpc user-agent header. Assert the header starts // with the // expected @@ -634,26 +635,28 @@ public void testSetCatalogShouldUpdateSessionOptions() throws Exception { 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)) { + 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 options = PRODUCER.getSessionOptions(); assertTrue(options.containsKey("catalog")); - String actualCatalog = options - .get("catalog") - .acceptVisitor( - new NoOpSessionOptionValueVisitor() { - @Override - public String visit(String value) { - return value; - } - }); + String actualCatalog = + options + .get("catalog") + .acceptVisitor( + new NoOpSessionOptionValueVisitor() { + @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 12f62f6729..45c2a96404 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 @@ -94,8 +94,10 @@ public final class MockFlightSqlProducer implements FlightSqlProducer { private final Map>> queryResults = new HashMap<>(); private final Map> selectResultProviders = new HashMap<>(); private final Map preparedStatements = new HashMap<>(); - private final Map> catalogQueriesResults = new HashMap<>(); - private final Map>> updateResultProviders = new HashMap<>(); + private final Map> catalogQueriesResults = + new HashMap<>(); + private final Map>> + updateResultProviders = new HashMap<>(); private final SqlInfoBuilder sqlInfoBuilder = new SqlInfoBuilder(); private final Map parameterSchemas = new HashMap<>(); private final Map>> expectedParameterValues = new HashMap<>(); @@ -131,8 +133,8 @@ public static ByteBuffer serializeSchema(final Schema schema) { /** * Registers a new {@link StatementType#SELECT} SQL query. * - * @param sqlCommand the SQL command under which to register the new query. - * @param schema the schema to use for the query result. + * @param sqlCommand the SQL command under which to register the new query. + * @param schema the schema to use for the query result. * @param resultProviders the result provider for this query. */ public void addSelectQuery( @@ -140,9 +142,10 @@ public void addSelectQuery( final Schema schema, final List> resultProviders) { final int providers = resultProviders.size(); - final List uuids = IntStream.range(0, providers) - .mapToObj(index -> new UUID(sqlCommand.hashCode(), Integer.hashCode(index))) - .collect(toList()); + final List uuids = + IntStream.range(0, providers) + .mapToObj(index -> new UUID(sqlCommand.hashCode(), Integer.hashCode(index))) + .collect(toList()); queryResults.put(sqlCommand, new SimpleImmutableEntry<>(schema, uuids)); IntStream.range(0, providers) .forEach( @@ -152,14 +155,15 @@ public void addSelectQuery( /** * Registers a new {@link StatementType#UPDATE} SQL query. * - * @param sqlCommand the SQL command. + * @param sqlCommand the SQL command. * @param updatedRows the number of rows affected. */ public void addUpdateQuery(final String sqlCommand, final long updatedRows) { addUpdateQuery( sqlCommand, (flightStream, putResultStreamListener) -> { - final DoPutUpdateResult result = DoPutUpdateResult.newBuilder().setRecordCount(updatedRows).build(); + final DoPutUpdateResult result = + DoPutUpdateResult.newBuilder().setRecordCount(updatedRows).build(); try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); final ArrowBuf buffer = allocator.buffer(result.getSerializedSize())) { buffer.writeBytes(result.toByteArray()); @@ -175,8 +179,7 @@ public void addUpdateQuery(final String sqlCommand, final long updatedRows) { /** * Adds a catalog query to the results. * - * @param message the {@link Message} corresponding to the catalog query - * request type to register. + * @param message the {@link Message} corresponding to the catalog query request type to register. * @param resultsProvider the results provider. */ public void addCatalogQuery( @@ -187,7 +190,7 @@ public void addCatalogQuery( /** * Registers a new {@link StatementType#UPDATE} SQL query. * - * @param sqlCommand the SQL command. + * @param sqlCommand the SQL command. * @param resultsProvider consumer for producing update results. */ void addUpdateQuery( @@ -211,11 +214,13 @@ public void createPreparedStatement( final CallContext callContext, final StreamListener listener) { try { - final ByteString preparedStatementHandle = copyFrom(randomUUID().toString().getBytes(StandardCharsets.UTF_8)); + final ByteString preparedStatementHandle = + copyFrom(randomUUID().toString().getBytes(StandardCharsets.UTF_8)); final String query = request.getQuery(); - final ActionCreatePreparedStatementResult.Builder resultBuilder = ActionCreatePreparedStatementResult.newBuilder() - .setPreparedStatementHandle(preparedStatementHandle); + final ActionCreatePreparedStatementResult.Builder resultBuilder = + ActionCreatePreparedStatementResult.newBuilder() + .setPreparedStatementHandle(preparedStatementHandle); final Entry> entry = queryResults.get(query); if (entry != null) { @@ -265,13 +270,15 @@ public FlightInfo getFlightInfoStatement( final CallContext callContext, final FlightDescriptor flightDescriptor) { final String query = commandStatementQuery.getQuery(); - final Entry> queryInfo = Preconditions.checkNotNull( - queryResults.get(query), format("Query not registered: <%s>.", query)); - final List endpoints = queryInfo.getValue().stream() - .map(TicketConversionUtils::getTicketBytesFromUuid) - .map(TicketConversionUtils::getTicketStatementQueryFromHandle) - .map(TicketConversionUtils::getEndpointFromMessage) - .collect(toList()); + final Entry> queryInfo = + Preconditions.checkNotNull( + queryResults.get(query), format("Query not registered: <%s>.", query)); + final List endpoints = + queryInfo.getValue().stream() + .map(TicketConversionUtils::getTicketBytesFromUuid) + .map(TicketConversionUtils::getTicketStatementQueryFromHandle) + .map(TicketConversionUtils::getEndpointFromMessage) + .collect(toList()); return FlightInfo.builder(queryInfo.getKey(), flightDescriptor, endpoints) .setAppMetadata("foo".getBytes(StandardCharsets.UTF_8)) .build(); @@ -282,18 +289,22 @@ public FlightInfo getFlightInfoPreparedStatement( final CommandPreparedStatementQuery commandPreparedStatementQuery, final CallContext callContext, final FlightDescriptor flightDescriptor) { - final ByteString preparedStatementHandle = commandPreparedStatementQuery.getPreparedStatementHandle(); - - final String query = Preconditions.checkNotNull( - preparedStatements.get(preparedStatementHandle), - format("No query registered under handle: <%s>.", preparedStatementHandle)); - final Entry> queryInfo = Preconditions.checkNotNull( - queryResults.get(query), format("Query not registered: <%s>.", query)); - final List endpoints = queryInfo.getValue().stream() - .map(TicketConversionUtils::getTicketBytesFromUuid) - .map(TicketConversionUtils::getCommandPreparedStatementQueryFromHandle) - .map(TicketConversionUtils::getEndpointFromMessage) - .collect(toList()); + final ByteString preparedStatementHandle = + commandPreparedStatementQuery.getPreparedStatementHandle(); + + final String query = + Preconditions.checkNotNull( + preparedStatements.get(preparedStatementHandle), + format("No query registered under handle: <%s>.", preparedStatementHandle)); + final Entry> queryInfo = + Preconditions.checkNotNull( + queryResults.get(query), format("Query not registered: <%s>.", query)); + final List endpoints = + queryInfo.getValue().stream() + .map(TicketConversionUtils::getTicketBytesFromUuid) + .map(TicketConversionUtils::getCommandPreparedStatementQueryFromHandle) + .map(TicketConversionUtils::getEndpointFromMessage) + .collect(toList()); return FlightInfo.builder(queryInfo.getKey(), flightDescriptor, endpoints) .setAppMetadata("foo".getBytes(StandardCharsets.UTF_8)) .build(); @@ -305,8 +316,9 @@ public SchemaResult getSchemaStatement( final CallContext callContext, final FlightDescriptor flightDescriptor) { final String query = commandStatementQuery.getQuery(); - final Entry> queryInfo = Preconditions.checkNotNull( - queryResults.get(query), format("Query not registered: <%s>.", query)); + final Entry> queryInfo = + Preconditions.checkNotNull( + queryResults.get(query), format("Query not registered: <%s>.", query)); return new SchemaResult(queryInfo.getKey()); } @@ -318,9 +330,9 @@ public void getStreamStatement( final ServerStreamListener serverStreamListener) { final UUID uuid = UUID.fromString(ticketStatementQuery.getStatementHandle().toStringUtf8()); Preconditions.checkNotNull( - selectResultProviders.get(uuid), - "No consumer was registered for the specified UUID: <%s>.", - uuid) + selectResultProviders.get(uuid), + "No consumer was registered for the specified UUID: <%s>.", + uuid) .accept(serverStreamListener); } @@ -329,11 +341,12 @@ public void getStreamPreparedStatement( final CommandPreparedStatementQuery commandPreparedStatementQuery, final CallContext callContext, final ServerStreamListener serverStreamListener) { - final UUID uuid = UUID.fromString(commandPreparedStatementQuery.getPreparedStatementHandle().toStringUtf8()); + final UUID uuid = + UUID.fromString(commandPreparedStatementQuery.getPreparedStatementHandle().toStringUtf8()); Preconditions.checkNotNull( - selectResultProviders.get(uuid), - "No consumer was registered for the specified UUID: <%s>.", - uuid) + selectResultProviders.get(uuid), + "No consumer was registered for the specified UUID: <%s>.", + uuid) .accept(serverStreamListener); } @@ -345,9 +358,10 @@ public Runnable acceptPutStatement( final StreamListener streamListener) { return () -> { final String query = commandStatementUpdate.getQuery(); - final BiConsumer> resultProvider = Preconditions.checkNotNull( - updateResultProviders.get(query), - format("No consumer found for query: <%s>.", query)); + final BiConsumer> resultProvider = + Preconditions.checkNotNull( + updateResultProviders.get(query), + format("No consumer found for query: <%s>.", query)); resultProvider.accept(flightStream, streamListener); }; } @@ -381,7 +395,8 @@ private boolean validateParameters( Object actual = root.getVector(paramIndex).getObject(i); boolean matches; if (expected.getClass().isArray()) { - matches = Arrays.equals((Object[]) expected, ((JsonStringArrayList) actual).toArray()); + matches = + Arrays.equals((Object[]) expected, ((JsonStringArrayList) actual).toArray()); } else { matches = Objects.equals(expected, actual); } @@ -415,13 +430,13 @@ public Runnable acceptPutPreparedStatementUpdate( final FlightStream flightStream, final StreamListener streamListener) { final ByteString handle = commandPreparedStatementUpdate.getPreparedStatementHandle(); - final String query = Preconditions.checkNotNull( - preparedStatements.get(handle), - format("No query registered under handle: <%s>.", handle)); + final String query = + Preconditions.checkNotNull( + preparedStatements.get(handle), + format("No query registered under handle: <%s>.", handle)); if (validateParameters(query, flightStream, streamListener)) { - return () -> { - }; + return () -> {}; } return acceptPutStatement( @@ -438,13 +453,13 @@ public Runnable acceptPutPreparedStatementQuery( final FlightStream flightStream, final StreamListener streamListener) { final ByteString handle = commandPreparedStatementQuery.getPreparedStatementHandle(); - final String query = Preconditions.checkNotNull( - preparedStatements.get(handle), - format("No query registered under handle: <%s>.", handle)); + final String query = + Preconditions.checkNotNull( + preparedStatements.get(handle), + format("No query registered under handle: <%s>.", handle)); if (validateParameters(query, flightStream, streamListener)) { - return () -> { - }; + return () -> {}; } return streamListener::onCompleted; @@ -629,8 +644,7 @@ public void doAction(CallContext context, Action action, StreamListener } /** - * Clear the `actionTypeCounter` map and restore to its default state. Intended - * to be used in + * Clear the `actionTypeCounter` map and restore to its default state. Intended to be used in * tests. */ public void clearActionTypeCounter() { @@ -644,8 +658,8 @@ public Map getActionTypeCounter() { private void getStreamCatalogFunctions( final Message ticket, final ServerStreamListener serverStreamListener) { Preconditions.checkNotNull( - catalogQueriesResults.get(ticket), - format("Query not registered for ticket: <%s>", ticket)) + catalogQueriesResults.get(ticket), + format("Query not registered for ticket: <%s>", ticket)) .accept(serverStreamListener); }