From fafd48c3196b4e6af0c09f633c6a3440419bf1f0 Mon Sep 17 00:00:00 2001 From: James Duong Date: Tue, 7 Nov 2023 08:40:38 -0800 Subject: [PATCH] GH-38576: [Java] Change JDBC driver to optionally preserve cookies and auth tokens when getting streams (#38580) ### Rationale for this change This change restores the original behavior of transmitting existing cookies and auth tokens when getting separate streams returned by getFlightInfo after adding support for multiple endpoints. These properties are now optional though. ### What changes are included in this PR? - Change the JDBC driver to add new properties "retainCookies" and "retainAuth" - These properties allow internally spawned connections for getting streams to use the cookies and bearer tokens from the original connection. - Add tests for validating defaults from ArrowFlightSqlClient.Builder ### Are these changes tested? Unit tests have been added. ### Are there any user-facing changes? Yes. There are now properties and they are documented. * Closes: #38576 Authored-by: James Duong Signed-off-by: David Li --- docs/source/java/flight_sql_jdbc_driver.rst | 11 ++ .../ClientIncomingAuthHeaderMiddleware.java | 2 +- .../driver/jdbc/ArrowFlightConnection.java | 2 + .../client/ArrowFlightSqlClientHandler.java | 117 +++++++++++--- .../ArrowFlightConnectionConfigImpl.java | 20 ++- .../driver/jdbc/ConnectionMutualTlsTest.java | 56 ++++--- .../arrow/driver/jdbc/ConnectionTest.java | 123 ++++++++------- .../jdbc/ConnectionTlsRootCertsTest.java | 68 +++++---- .../arrow/driver/jdbc/ConnectionTlsTest.java | 56 +++---- .../driver/jdbc/FlightServerTestRule.java | 5 +- .../arrow/driver/jdbc/ResultSetTest.java | 15 +- ...rrowFlightSqlClientHandlerBuilderTest.java | 143 ++++++++++++++++++ 12 files changed, 448 insertions(+), 170 deletions(-) create mode 100644 java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerBuilderTest.java diff --git a/docs/source/java/flight_sql_jdbc_driver.rst b/docs/source/java/flight_sql_jdbc_driver.rst index 34ccfea47f9e3..0ace2185983a9 100644 --- a/docs/source/java/flight_sql_jdbc_driver.rst +++ b/docs/source/java/flight_sql_jdbc_driver.rst @@ -141,6 +141,17 @@ case-sensitive. The supported parameters are: - true - When TLS is enabled, whether to use the system certificate store + * - retainCookies + - true + - Whether to use cookies from the initial connection in subsequent + internal connections when retrieving streams from separate endpoints. + + * - retainAuth + - true + - Whether to use bearer tokens obtained from the initial connection + in subsequent internal connections used for retrieving streams + from separate endpoints. + Note that URI values must be URI-encoded if they contain characters such as !, @, $, etc. diff --git a/java/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/ClientIncomingAuthHeaderMiddleware.java b/java/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/ClientIncomingAuthHeaderMiddleware.java index be5f3f54d326c..7bb55d145d104 100644 --- a/java/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/ClientIncomingAuthHeaderMiddleware.java +++ b/java/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/ClientIncomingAuthHeaderMiddleware.java @@ -34,7 +34,7 @@ public class ClientIncomingAuthHeaderMiddleware implements FlightClientMiddlewar */ public static class Factory implements FlightClientMiddleware.Factory { private final ClientHeaderHandler headerHandler; - private CredentialCallOption credentialCallOption; + private CredentialCallOption credentialCallOption = null; /** * Construct a factory with the given header handler. diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java index fdbb9381c0a55..ad19c616ff29a 100644 --- a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java @@ -109,6 +109,8 @@ private static ArrowFlightSqlClientHandler createNewClientHandler( .withDisableCertificateVerification(config.getDisableCertificateVerification()) .withToken(config.getToken()) .withCallOptions(config.toCallOption()) + .withRetainCookies(config.retainCookies()) + .withRetainAuth(config.retainAuth()) .build(); } catch (final SQLException e) { try { diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java index 75e80d45dc669..54fd17853c00b 100644 --- a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java @@ -49,6 +49,7 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; +import org.apache.arrow.util.VisibleForTesting; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.calcite.avatica.Meta.StatementType; @@ -425,21 +426,59 @@ public static final class Builder { private final Set options = new HashSet<>(); private String host; private int port; - private String username; - private String password; - private String trustStorePath; - private String trustStorePassword; - private String token; - private boolean useEncryption; - private boolean disableCertificateVerification; - private boolean useSystemTrustStore; - private String tlsRootCertificatesPath; - private String clientCertificatePath; - private String clientKeyPath; + + @VisibleForTesting + String username; + + @VisibleForTesting + String password; + + @VisibleForTesting + String trustStorePath; + + @VisibleForTesting + String trustStorePassword; + + @VisibleForTesting + String token; + + @VisibleForTesting + boolean useEncryption = true; + + @VisibleForTesting + boolean disableCertificateVerification; + + @VisibleForTesting + boolean useSystemTrustStore = true; + + @VisibleForTesting + String tlsRootCertificatesPath; + + @VisibleForTesting + String clientCertificatePath; + + @VisibleForTesting + String clientKeyPath; + + @VisibleForTesting private BufferAllocator allocator; - public Builder() { + @VisibleForTesting + boolean retainCookies = true; + + @VisibleForTesting + boolean retainAuth = true; + + // These two middlewares are for internal use within build() and should not be exposed by builder APIs. + // Note that these middlewares may not necessarily be registered. + @VisibleForTesting + ClientIncomingAuthHeaderMiddleware.Factory authFactory + = new ClientIncomingAuthHeaderMiddleware.Factory(new ClientBearerHeaderHandler()); + @VisibleForTesting + ClientCookieMiddleware.Factory cookieFactory = new ClientCookieMiddleware.Factory(); + + public Builder() { } /** @@ -447,7 +486,8 @@ public Builder() { * * @param original The builder to base this copy off of. */ - private Builder(Builder original) { + @VisibleForTesting + Builder(Builder original) { this.middlewareFactories.addAll(original.middlewareFactories); this.options.addAll(original.options); this.host = original.host; @@ -464,6 +504,14 @@ private Builder(Builder original) { this.clientCertificatePath = original.clientCertificatePath; this.clientKeyPath = original.clientKeyPath; this.allocator = original.allocator; + + if (original.retainCookies) { + this.cookieFactory = original.cookieFactory; + } + + if (original.retainAuth) { + this.authFactory = original.authFactory; + } } /** @@ -622,6 +670,28 @@ public Builder withBufferAllocator(final BufferAllocator allocator) { return this; } + /** + * Indicates if cookies should be re-used by connections spawned for getStreams() calls. + * @param retainCookies The flag indicating if cookies should be re-used. + * @return this builder instance. + */ + public Builder withRetainCookies(boolean retainCookies) { + this.retainCookies = retainCookies; + return this; + } + + /** + * Indicates if bearer tokens negotiated should be re-used by connections + * spawned for getStreams() calls. + * + * @param retainAuth The flag indicating if auth tokens should be re-used. + * @return this builder instance. + */ + public Builder withRetainAuth(boolean retainAuth) { + this.retainAuth = retainAuth; + return this; + } + /** * Adds the provided {@code factories} to the list of {@link #middlewareFactories} of this handler. * @@ -675,13 +745,11 @@ public ArrowFlightSqlClientHandler build() throws SQLException { // Copy middlewares so that the build method doesn't change the state of the builder fields itself. Set buildTimeMiddlewareFactories = new HashSet<>(this.middlewareFactories); FlightClient client = null; + boolean isUsingUserPasswordAuth = username != null && token == null; try { - ClientIncomingAuthHeaderMiddleware.Factory authFactory = null; // Token should take priority since some apps pass in a username/password even when a token is provided - if (username != null && token == null) { - authFactory = - new ClientIncomingAuthHeaderMiddleware.Factory(new ClientBearerHeaderHandler()); + if (isUsingUserPasswordAuth) { buildTimeMiddlewareFactories.add(authFactory); } final FlightClient.Builder clientBuilder = FlightClient.builder().allocator(allocator); @@ -722,10 +790,17 @@ public ArrowFlightSqlClientHandler build() throws SQLException { client = clientBuilder.build(); final ArrayList credentialOptions = new ArrayList<>(); - if (authFactory != null) { - credentialOptions.add( - ClientAuthenticationUtils.getAuthenticate( - client, username, password, authFactory, options.toArray(new CallOption[0]))); + 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 getStream(). + if (authFactory.getCredentialCallOption() != null) { + credentialOptions.add(authFactory.getCredentialCallOption()); + } else { + // Otherwise do the handshake and get the token if possible. + credentialOptions.add( + ClientAuthenticationUtils.getAuthenticate( + client, username, password, authFactory, options.toArray(new CallOption[0]))); + } } else if (token != null) { credentialOptions.add( ClientAuthenticationUtils.getAuthenticate( diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java index 59118e1d6f788..6237a8b58d68a 100644 --- a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java @@ -143,6 +143,22 @@ public int threadPoolSize() { return ArrowFlightConnectionProperty.THREAD_POOL_SIZE.getInteger(properties); } + /** + * Indicates if sub-connections created for stream retrieval + * should reuse cookies from the main connection. + */ + public boolean retainCookies() { + return ArrowFlightConnectionProperty.RETAIN_COOKIES.getBoolean(properties); + } + + /** + * Indicates if sub-connections created for stream retrieval + * should reuse bearer tokens created from the main connection. + */ + public boolean retainAuth() { + return ArrowFlightConnectionProperty.RETAIN_AUTH.getBoolean(properties); + } + /** * Gets the {@link CallOption}s from this {@link ConnectionConfig}. * @@ -191,7 +207,9 @@ public enum ArrowFlightConnectionProperty implements ConnectionProperty { CLIENT_CERTIFICATE("clientCertificate", null, Type.STRING, false), CLIENT_KEY("clientKey", null, Type.STRING, false), THREAD_POOL_SIZE("threadPoolSize", 1, Type.NUMBER, false), - TOKEN("token", null, Type.STRING, false); + TOKEN("token", null, Type.STRING, false), + RETAIN_COOKIES("retainCookies", true, Type.BOOLEAN, false), + RETAIN_AUTH("retainAuth", true, Type.BOOLEAN, false); private final String camelName; private final Object defaultValue; diff --git a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionMutualTlsTest.java b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionMutualTlsTest.java index 783e0c41e9269..927b3e426c6ba 100644 --- a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionMutualTlsTest.java +++ b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionMutualTlsTest.java @@ -140,17 +140,21 @@ public void testGetEncryptedClientWithBadMTlsCertPath() { final UsernamePasswordCredentials credentials = new UsernamePasswordCredentials( userTest, passTest); - assertThrows(SQLException.class, () -> new ArrowFlightSqlClientHandler.Builder() - .withHost(FLIGHT_SERVER_TEST_RULE.getHost()) - .withPort(FLIGHT_SERVER_TEST_RULE.getPort()) - .withUsername(credentials.getUserName()) - .withPassword(credentials.getPassword()) - .withTlsRootCertificates(tlsRootCertsPath) - .withClientCertificate(badClientMTlsCertPath) - .withClientKey(clientMTlsKeyPath) - .withBufferAllocator(allocator) - .withEncryption(true) - .build()); + assertThrows(SQLException.class, () -> { + try (ArrowFlightSqlClientHandler handler = new ArrowFlightSqlClientHandler.Builder() + .withHost(FLIGHT_SERVER_TEST_RULE.getHost()) + .withPort(FLIGHT_SERVER_TEST_RULE.getPort()) + .withUsername(credentials.getUserName()) + .withPassword(credentials.getPassword()) + .withTlsRootCertificates(tlsRootCertsPath) + .withClientCertificate(badClientMTlsCertPath) + .withClientKey(clientMTlsKeyPath) + .withBufferAllocator(allocator) + .withEncryption(true) + .build()) { + Assert.fail(); + } + }); } /** @@ -162,17 +166,21 @@ public void testGetEncryptedClientWithBadMTlsKeyPath() { final UsernamePasswordCredentials credentials = new UsernamePasswordCredentials( userTest, passTest); - assertThrows(SQLException.class, () -> new ArrowFlightSqlClientHandler.Builder() - .withHost(FLIGHT_SERVER_TEST_RULE.getHost()) - .withPort(FLIGHT_SERVER_TEST_RULE.getPort()) - .withUsername(credentials.getUserName()) - .withPassword(credentials.getPassword()) - .withTlsRootCertificates(tlsRootCertsPath) - .withClientCertificate(clientMTlsCertPath) - .withClientKey(badClientMTlsKeyPath) - .withBufferAllocator(allocator) - .withEncryption(true) - .build()); + assertThrows(SQLException.class, () -> { + try (ArrowFlightSqlClientHandler handler = new ArrowFlightSqlClientHandler.Builder() + .withHost(FLIGHT_SERVER_TEST_RULE.getHost()) + .withPort(FLIGHT_SERVER_TEST_RULE.getPort()) + .withUsername(credentials.getUserName()) + .withPassword(credentials.getPassword()) + .withTlsRootCertificates(tlsRootCertsPath) + .withClientCertificate(clientMTlsCertPath) + .withClientKey(badClientMTlsKeyPath) + .withBufferAllocator(allocator) + .withEncryption(true) + .build()) { + Assert.fail(); + } + }); } /** @@ -222,7 +230,7 @@ public void testGetEncryptedConnectionWithValidCredentialsAndTlsRootsPath() thro final ArrowFlightJdbcDataSource dataSource = ArrowFlightJdbcDataSource.createNewDataSource(properties); try (final Connection connection = dataSource.getConnection()) { - assert connection.isValid(300); + Assert.assertTrue(connection.isValid(300)); } } @@ -245,7 +253,7 @@ public void testGetNonAuthenticatedEncryptedConnection() throws Exception { final ArrowFlightJdbcDataSource dataSource = ArrowFlightJdbcDataSource.createNewDataSource(properties); try (final Connection connection = dataSource.getConnection()) { - assert connection.isValid(300); + Assert.assertTrue(connection.isValid(300)); } } diff --git a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java index bec0ff1e59752..357506b3d177c 100644 --- a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java +++ b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java @@ -20,6 +20,7 @@ import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; import java.net.URISyntaxException; import java.sql.Connection; @@ -100,7 +101,7 @@ public void testUnencryptedConnectionShouldOpenSuccessfullyWhenProvidedValidCred try (Connection connection = DriverManager.getConnection( "jdbc:arrow-flight-sql://" + FLIGHT_SERVER_TEST_RULE.getHost() + ":" + FLIGHT_SERVER_TEST_RULE.getPort(), properties)) { - assert connection.isValid(300); + Assert.assertTrue(connection.isValid(300)); } } @@ -122,11 +123,13 @@ public void testTokenOverridesUsernameAndPasswordAuth() { properties.put(ArrowFlightConnectionProperty.TOKEN.camelName(), "token"); properties.put("useEncryption", false); - SQLException e = assertThrows(SQLException.class, () -> - DriverManager.getConnection( - "jdbc:arrow-flight-sql://" + FLIGHT_SERVER_TEST_RULE.getHost() + ":" + - FLIGHT_SERVER_TEST_RULE.getPort(), - properties)); + SQLException e = assertThrows(SQLException.class, () -> { + try (Connection conn = DriverManager.getConnection( + "jdbc:arrow-flight-sql://" + FLIGHT_SERVER_TEST_RULE.getHost() + ":" + FLIGHT_SERVER_TEST_RULE.getPort(), + properties)) { + Assert.fail(); + } + }); assertTrue(e.getMessage().contains("UNAUTHENTICATED")); } @@ -145,7 +148,9 @@ public void testUnencryptedConnectionWithEmptyHost() properties.put("password", passTest); final String invalidUrl = "jdbc:arrow-flight-sql://"; - DriverManager.getConnection(invalidUrl, properties); + try (Connection conn = DriverManager.getConnection(invalidUrl, properties)) { + Assert.fail("Expected SQLException."); + } } /** @@ -161,6 +166,7 @@ public void testGetBasicClientAuthenticatedShouldOpenConnection() new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_RULE.getHost()) .withPort(FLIGHT_SERVER_TEST_RULE.getPort()) + .withEncryption(false) .withUsername(userTest) .withPassword(passTest) .withBufferAllocator(allocator) @@ -191,7 +197,9 @@ public void testUnencryptedConnectionProvidingInvalidPort() final String invalidUrl = "jdbc:arrow-flight-sql://" + FLIGHT_SERVER_TEST_RULE.getHost() + ":" + 65537; - DriverManager.getConnection(invalidUrl, properties); + try (Connection conn = DriverManager.getConnection(invalidUrl, properties)) { + fail("Expected SQLException"); + } } /** @@ -206,6 +214,7 @@ public void testGetBasicClientNoAuthShouldOpenConnection() throws Exception { new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_RULE.getHost()) .withBufferAllocator(allocator) + .withEncryption(false) .build()) { assertNotNull(client); } @@ -228,7 +237,7 @@ public void testUnencryptedConnectionShouldOpenSuccessfullyWithoutAuthentication false); try (Connection connection = DriverManager .getConnection("jdbc:arrow-flight-sql://localhost:32010", properties)) { - assert connection.isValid(300); + Assert.assertTrue(connection.isValid(300)); } } @@ -271,14 +280,14 @@ public void testTLSConnectionPropertyFalseCorrectCastUrlWithDriverManager() thro final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); - Connection connection = DriverManager.getConnection( + try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&useEncryption=false", FLIGHT_SERVER_TEST_RULE.getPort(), userTest, - passTest)); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + passTest))) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -301,13 +310,13 @@ public void testTLSConnectionPropertyFalseCorrectCastUrlAndPropertiesUsingSetPro passTest); properties.setProperty(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), "false"); - Connection connection = DriverManager.getConnection( + try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_RULE.getPort()), - properties); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + properties)) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -329,13 +338,13 @@ public void testTLSConnectionPropertyFalseCorrectCastUrlAndPropertiesUsingPutWit passTest); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), false); - Connection connection = DriverManager.getConnection( + try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_RULE.getPort()), - properties); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + properties)) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -350,14 +359,14 @@ public void testTLSConnectionPropertyFalseIntegerCorrectCastUrlWithDriverManager final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); - Connection connection = DriverManager.getConnection( + try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&useEncryption=0", FLIGHT_SERVER_TEST_RULE.getPort(), userTest, - passTest)); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + passTest))) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -380,13 +389,13 @@ public void testTLSConnectionPropertyFalseIntegerCorrectCastUrlAndPropertiesUsin passTest); properties.setProperty(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), "0"); - Connection connection = DriverManager.getConnection( + try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_RULE.getPort()), - properties); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + properties)) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -409,13 +418,13 @@ public void testTLSConnectionPropertyFalseIntegerCorrectCastUrlAndPropertiesUsin passTest); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), 0); - Connection connection = DriverManager.getConnection( + try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_RULE.getPort()), - properties); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + properties)) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -430,15 +439,15 @@ public void testThreadPoolSizeConnectionPropertyCorrectCastUrlWithDriverManager( final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); - Connection connection = DriverManager.getConnection( + try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&threadPoolSize=1&useEncryption=%s", FLIGHT_SERVER_TEST_RULE.getPort(), userTest, passTest, - false)); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + false))) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -462,13 +471,13 @@ public void testThreadPoolSizeConnectionPropertyCorrectCastUrlAndPropertiesUsing properties.setProperty(ArrowFlightConnectionProperty.THREAD_POOL_SIZE.camelName(), "1"); properties.put("useEncryption", false); - Connection connection = DriverManager.getConnection( + try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_RULE.getPort()), - properties); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + properties)) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -492,13 +501,13 @@ public void testThreadPoolSizeConnectionPropertyCorrectCastUrlAndPropertiesUsing properties.put(ArrowFlightConnectionProperty.THREAD_POOL_SIZE.camelName(), 1); properties.put("useEncryption", false); - Connection connection = DriverManager.getConnection( + try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_RULE.getPort()), - properties); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + properties)) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -513,15 +522,15 @@ public void testPasswordConnectionPropertyIntegerCorrectCastUrlWithDriverManager final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); - Connection connection = DriverManager.getConnection( + try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&useEncryption=%s", FLIGHT_SERVER_TEST_RULE.getPort(), userTest, passTest, - false)); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + false))) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -544,13 +553,13 @@ public void testPasswordConnectionPropertyIntegerCorrectCastUrlAndPropertiesUsin passTest); properties.put("useEncryption", false); - Connection connection = DriverManager.getConnection( + try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_RULE.getPort()), - properties); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + properties)) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -573,12 +582,12 @@ public void testPasswordConnectionPropertyIntegerCorrectCastUrlAndPropertiesUsin passTest); properties.put("useEncryption", false); - Connection connection = DriverManager.getConnection( + try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_RULE.getPort()), - properties); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + properties)) { + Assert.assertTrue(connection.isValid(0)); + } } } diff --git a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTlsRootCertsTest.java b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTlsRootCertsTest.java index 333c474824e52..5579cf0cf5f54 100644 --- a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTlsRootCertsTest.java +++ b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTlsRootCertsTest.java @@ -118,12 +118,16 @@ public void testGetEncryptedClientAuthenticated() throws Exception { */ @Test public void testGetEncryptedClientWithNoCertificateOnKeyStore() { - assertThrows(SQLException.class, () -> new ArrowFlightSqlClientHandler.Builder() - .withHost(FLIGHT_SERVER_TEST_RULE.getHost()) - .withTlsRootCertificates(badTlsRootCertsPath) - .withBufferAllocator(allocator) - .withEncryption(true) - .build()); + assertThrows(SQLException.class, () -> { + try (ArrowFlightSqlClientHandler handler = new ArrowFlightSqlClientHandler.Builder() + .withHost(FLIGHT_SERVER_TEST_RULE.getHost()) + .withTlsRootCertificates(badTlsRootCertsPath) + .withBufferAllocator(allocator) + .withEncryption(true) + .build()) { + Assert.fail(); + } + }); } /** @@ -167,7 +171,7 @@ public void testGetEncryptedConnectionWithValidCredentialsAndTlsRootsPath() thro final ArrowFlightJdbcDataSource dataSource = ArrowFlightJdbcDataSource.createNewDataSource(properties); try (final Connection connection = dataSource.getConnection()) { - assert connection.isValid(300); + Assert.assertTrue(connection.isValid(300)); } } @@ -188,7 +192,7 @@ public void testGetNonAuthenticatedEncryptedConnection() throws Exception { final ArrowFlightJdbcDataSource dataSource = ArrowFlightJdbcDataSource.createNewDataSource(properties); try (final Connection connection = dataSource.getConnection()) { - assert connection.isValid(300); + Assert.assertTrue(connection.isValid(300)); } } @@ -203,7 +207,7 @@ public void testTLSConnectionPropertyTrueCorrectCastUrlWithDriverManager() throw final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); - final Connection connection = DriverManager.getConnection( + try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s" + "&useEncryption=true&%s=%s", @@ -211,9 +215,9 @@ public void testTLSConnectionPropertyTrueCorrectCastUrlWithDriverManager() throw userTest, passTest, ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), - URLEncoder.encode(tlsRootCertsPath, "UTF-8"))); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + URLEncoder.encode(tlsRootCertsPath, "UTF-8")))) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -235,13 +239,13 @@ public void testTLSConnectionPropertyTrueCorrectCastUrlAndPropertiesUsingSetProp properties.setProperty(ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), tlsRootCertsPath); properties.setProperty(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), "true"); - final Connection connection = DriverManager.getConnection( + try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_RULE.getPort()), - properties); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + properties)) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -263,13 +267,13 @@ public void testTLSConnectionPropertyTrueCorrectCastUrlAndPropertiesUsingPutWith properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), true); properties.put(ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), tlsRootCertsPath); - final Connection connection = DriverManager.getConnection( + try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_RULE.getPort()), - properties); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + properties)) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -284,7 +288,7 @@ public void testTLSConnectionPropertyTrueIntegerCorrectCastUrlWithDriverManager( final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); - final Connection connection = DriverManager.getConnection( + try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s" + "&useEncryption=1&useSystemTrustStore=0&%s=%s", @@ -292,9 +296,9 @@ public void testTLSConnectionPropertyTrueIntegerCorrectCastUrlWithDriverManager( userTest, passTest, ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), - URLEncoder.encode(tlsRootCertsPath, "UTF-8"))); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + URLEncoder.encode(tlsRootCertsPath, "UTF-8")))) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -316,11 +320,11 @@ public void testTLSConnectionPropertyTrueIntegerCorrectCastUrlAndPropertiesUsing properties.setProperty(ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), tlsRootCertsPath); properties.setProperty(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), "1"); - final Connection connection = DriverManager.getConnection( + try (final Connection connection = DriverManager.getConnection( String.format("jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_RULE.getPort()), - properties); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + properties)) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -342,11 +346,11 @@ public void testTLSConnectionPropertyTrueIntegerCorrectCastUrlAndPropertiesUsing properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), 1); properties.put(ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), tlsRootCertsPath); - final Connection connection = DriverManager.getConnection( + try (final Connection connection = DriverManager.getConnection( String.format("jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_RULE.getPort()), - properties); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + properties)) { + Assert.assertTrue(connection.isValid(0)); + } } } diff --git a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTlsTest.java b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTlsTest.java index 95d591766a836..7e160f3f0c385 100644 --- a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTlsTest.java +++ b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTlsTest.java @@ -127,6 +127,7 @@ public void testGetEncryptedClientAuthenticated() throws Exception { new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_RULE.getHost()) .withPort(FLIGHT_SERVER_TEST_RULE.getPort()) + .withSystemTrustStore(false) .withUsername(credentials.getUserName()) .withPassword(credentials.getPassword()) .withTrustStorePath(trustStorePath) @@ -153,6 +154,7 @@ public void testGetEncryptedClientWithNoCertificateOnKeyStore() throws Exception .withHost(FLIGHT_SERVER_TEST_RULE.getHost()) .withTrustStorePath(noCertificateKeyStorePath) .withTrustStorePassword(noCertificateKeyStorePassword) + .withSystemTrustStore(false) .withBufferAllocator(allocator) .withEncryption(true) .build()) { @@ -170,6 +172,7 @@ public void testGetNonAuthenticatedEncryptedClientNoAuth() throws Exception { try (ArrowFlightSqlClientHandler client = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_RULE.getHost()) + .withSystemTrustStore(false) .withTrustStorePath(trustStorePath) .withTrustStorePassword(trustStorePass) .withBufferAllocator(allocator) @@ -192,6 +195,7 @@ public void testGetEncryptedClientWithKeyStoreBadPasswordAndNoAuth() throws Exce try (ArrowFlightSqlClientHandler ignored = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_RULE.getHost()) + .withSystemTrustStore(false) .withTrustStorePath(trustStorePath) .withTrustStorePassword(keyStoreBadPassword) .withBufferAllocator(allocator) @@ -225,7 +229,7 @@ public void testGetEncryptedConnectionWithValidCredentialsAndKeyStore() throws E final ArrowFlightJdbcDataSource dataSource = ArrowFlightJdbcDataSource.createNewDataSource(properties); try (final Connection connection = dataSource.getConnection()) { - assert connection.isValid(300); + Assert.assertTrue(connection.isValid(300)); } } @@ -276,7 +280,7 @@ public void testGetNonAuthenticatedEncryptedConnection() throws Exception { final ArrowFlightJdbcDataSource dataSource = ArrowFlightJdbcDataSource.createNewDataSource(properties); try (final Connection connection = dataSource.getConnection()) { - assert connection.isValid(300); + Assert.assertTrue(connection.isValid(300)); } } @@ -291,7 +295,7 @@ public void testTLSConnectionPropertyTrueCorrectCastUrlWithDriverManager() throw final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); - final Connection connection = DriverManager.getConnection( + try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s" + "&useEncryption=true&useSystemTrustStore=false&%s=%s&%s=%s", @@ -301,9 +305,9 @@ public void testTLSConnectionPropertyTrueCorrectCastUrlWithDriverManager() throw ArrowFlightConnectionProperty.TRUST_STORE.camelName(), URLEncoder.encode(trustStorePath, "UTF-8"), ArrowFlightConnectionProperty.TRUST_STORE_PASSWORD.camelName(), - URLEncoder.encode(trustStorePass, "UTF-8"))); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + URLEncoder.encode(trustStorePass, "UTF-8")))) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -327,13 +331,13 @@ public void testTLSConnectionPropertyTrueCorrectCastUrlAndPropertiesUsingSetProp properties.setProperty(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), "true"); properties.setProperty(ArrowFlightConnectionProperty.USE_SYSTEM_TRUST_STORE.camelName(), "false"); - final Connection connection = DriverManager.getConnection( + try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_RULE.getPort()), - properties); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + properties)) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -357,13 +361,13 @@ public void testTLSConnectionPropertyTrueCorrectCastUrlAndPropertiesUsingPutWith properties.put(ArrowFlightConnectionProperty.TRUST_STORE.camelName(), trustStorePath); properties.put(ArrowFlightConnectionProperty.TRUST_STORE_PASSWORD.camelName(), trustStorePass); - final Connection connection = DriverManager.getConnection( + try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_RULE.getPort()), - properties); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + properties)) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -378,7 +382,7 @@ public void testTLSConnectionPropertyTrueIntegerCorrectCastUrlWithDriverManager( final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); - final Connection connection = DriverManager.getConnection( + try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s" + "&useEncryption=1&useSystemTrustStore=0&%s=%s&%s=%s", @@ -388,9 +392,9 @@ public void testTLSConnectionPropertyTrueIntegerCorrectCastUrlWithDriverManager( ArrowFlightConnectionProperty.TRUST_STORE.camelName(), URLEncoder.encode(trustStorePath, "UTF-8"), ArrowFlightConnectionProperty.TRUST_STORE_PASSWORD.camelName(), - URLEncoder.encode(trustStorePass, "UTF-8"))); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + URLEncoder.encode(trustStorePass, "UTF-8")))) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -414,11 +418,11 @@ public void testTLSConnectionPropertyTrueIntegerCorrectCastUrlAndPropertiesUsing properties.setProperty(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), "1"); properties.setProperty(ArrowFlightConnectionProperty.USE_SYSTEM_TRUST_STORE.camelName(), "0"); - final Connection connection = DriverManager.getConnection( + try (final Connection connection = DriverManager.getConnection( String.format("jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_RULE.getPort()), - properties); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + properties)) { + Assert.assertTrue(connection.isValid(0)); + } } /** @@ -442,11 +446,11 @@ public void testTLSConnectionPropertyTrueIntegerCorrectCastUrlAndPropertiesUsing properties.put(ArrowFlightConnectionProperty.TRUST_STORE.camelName(), trustStorePath); properties.put(ArrowFlightConnectionProperty.TRUST_STORE_PASSWORD.camelName(), trustStorePass); - final Connection connection = DriverManager.getConnection( + try (final Connection connection = DriverManager.getConnection( String.format("jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_RULE.getPort()), - properties); - Assert.assertTrue(connection.isValid(0)); - connection.close(); + properties)) { + Assert.assertTrue(connection.isValid(0)); + } } } diff --git a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestRule.java b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestRule.java index df7cbea56ee2f..39eb0a29866f1 100644 --- a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestRule.java +++ b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestRule.java @@ -55,6 +55,9 @@ * and interact with it. */ public class FlightServerTestRule implements TestRule, AutoCloseable { + public static final String DEFAULT_USER = "flight-test-user"; + public static final String DEFAULT_PASSWORD = "flight-test-password"; + private static final Logger LOGGER = LoggerFactory.getLogger(FlightServerTestRule.class); private final Properties properties; @@ -92,7 +95,7 @@ private FlightServerTestRule(final Properties properties, public static FlightServerTestRule createStandardTestRule(final FlightSqlProducer producer) { UserPasswordAuthentication authentication = new UserPasswordAuthentication.Builder() - .user("flight-test-user", "flight-test-password") + .user(DEFAULT_USER, DEFAULT_PASSWORD) .build(); return new Builder() diff --git a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java index 5b9e269fb3cc8..52910812fb4fb 100644 --- a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java +++ b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java @@ -237,16 +237,17 @@ public void testShouldCloseStatementWhenIsCloseOnCompletion() throws Exception { */ @Test public void testShouldCloseStatementWhenIsCloseOnCompletionWithMaxRowsLimit() throws Exception { - Statement statement = connection.createStatement(); - ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD); + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { - final long maxRowsLimit = 3; - statement.setLargeMaxRows(maxRowsLimit); - statement.closeOnCompletion(); + final long maxRowsLimit = 3; + statement.setLargeMaxRows(maxRowsLimit); + statement.closeOnCompletion(); - resultSetNextUntilDone(resultSet); + resultSetNextUntilDone(resultSet); - collector.checkThat(statement.isClosed(), is(true)); + collector.checkThat(statement.isClosed(), is(true)); + } } /** diff --git a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerBuilderTest.java b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerBuilderTest.java new file mode 100644 index 0000000000000..6565a85ddf99f --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerBuilderTest.java @@ -0,0 +1,143 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.arrow.driver.jdbc.client; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotSame; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; + +import org.apache.arrow.driver.jdbc.FlightServerTestRule; +import org.apache.arrow.driver.jdbc.utils.CoreMockedSqlProducers; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.ClassRule; +import org.junit.Test; + +/** + * Test the behavior of ArrowFlightSqlClientHandler.Builder + */ +public class ArrowFlightSqlClientHandlerBuilderTest { + @ClassRule + public static final FlightServerTestRule FLIGHT_SERVER_TEST_RULE = FlightServerTestRule + .createStandardTestRule(CoreMockedSqlProducers.getLegacyProducer()); + + private static BufferAllocator allocator; + + @BeforeClass + public static void setup() { + allocator = new RootAllocator(Long.MAX_VALUE); + } + + @AfterClass + public static void tearDown() { + allocator.close(); + } + + @Test + public void testRetainCookiesOnAuthOff() throws Exception { + // Arrange + final ArrowFlightSqlClientHandler.Builder rootBuilder = new ArrowFlightSqlClientHandler.Builder() + .withHost(FLIGHT_SERVER_TEST_RULE.getHost()) + .withPort(FLIGHT_SERVER_TEST_RULE.getPort()) + .withBufferAllocator(allocator) + .withUsername(FlightServerTestRule.DEFAULT_USER) + .withPassword(FlightServerTestRule.DEFAULT_PASSWORD) + .withEncryption(false) + .withRetainCookies(true) + .withRetainAuth(false); + + try (ArrowFlightSqlClientHandler rootHandler = rootBuilder.build()) { + // Act + final ArrowFlightSqlClientHandler.Builder testBuilder = new ArrowFlightSqlClientHandler.Builder(rootBuilder); + + // Assert + assertSame(rootBuilder.cookieFactory, testBuilder.cookieFactory); + assertNotSame(rootBuilder.authFactory, testBuilder.authFactory); + } + } + + @Test + public void testRetainCookiesOffAuthOff() throws Exception { + // Arrange + final ArrowFlightSqlClientHandler.Builder rootBuilder = new ArrowFlightSqlClientHandler.Builder() + .withHost(FLIGHT_SERVER_TEST_RULE.getHost()) + .withPort(FLIGHT_SERVER_TEST_RULE.getPort()) + .withBufferAllocator(allocator) + .withUsername(FlightServerTestRule.DEFAULT_USER) + .withPassword(FlightServerTestRule.DEFAULT_PASSWORD) + .withEncryption(false) + .withRetainCookies(false) + .withRetainAuth(false); + + try (ArrowFlightSqlClientHandler rootHandler = rootBuilder.build()) { + // Act + final ArrowFlightSqlClientHandler.Builder testBuilder = new ArrowFlightSqlClientHandler.Builder(rootBuilder); + + // Assert + assertNotSame(rootBuilder.cookieFactory, testBuilder.cookieFactory); + assertNotSame(rootBuilder.authFactory, testBuilder.authFactory); + } + } + + @Test + public void testRetainCookiesOnAuthOn() throws Exception { + // Arrange + final ArrowFlightSqlClientHandler.Builder rootBuilder = new ArrowFlightSqlClientHandler.Builder() + .withHost(FLIGHT_SERVER_TEST_RULE.getHost()) + .withPort(FLIGHT_SERVER_TEST_RULE.getPort()) + .withBufferAllocator(allocator) + .withUsername(FlightServerTestRule.DEFAULT_USER) + .withPassword(FlightServerTestRule.DEFAULT_PASSWORD) + .withEncryption(false) + .withRetainCookies(true) + .withRetainAuth(true); + + try (ArrowFlightSqlClientHandler rootHandler = rootBuilder.build()) { + // Act + final ArrowFlightSqlClientHandler.Builder testBuilder = new ArrowFlightSqlClientHandler.Builder(rootBuilder); + + // Assert + assertSame(rootBuilder.cookieFactory, testBuilder.cookieFactory); + assertSame(rootBuilder.authFactory, testBuilder.authFactory); + } + } + + @Test + public void testDefaults() { + final ArrowFlightSqlClientHandler.Builder builder = new ArrowFlightSqlClientHandler.Builder(); + + // Validate all non-mandatory fields against defaults in ArrowFlightConnectionProperty. + assertNull(builder.username); + assertNull(builder.password); + assertTrue(builder.useEncryption); + assertFalse(builder.disableCertificateVerification); + assertNull(builder.trustStorePath); + assertNull(builder.trustStorePassword); + assertTrue(builder.useSystemTrustStore); + assertNull(builder.token); + assertTrue(builder.retainAuth); + assertTrue(builder.retainCookies); + assertNull(builder.tlsRootCertificatesPath); + assertNull(builder.clientCertificatePath); + assertNull(builder.clientKeyPath); + } +}