From a1980feb033834152299009fb77e8d0b8bf9b0ee Mon Sep 17 00:00:00 2001 From: alexradzin Date: Mon, 17 Jun 2024 14:19:01 +0300 Subject: [PATCH 1/3] FIR-33849 close all open connections when process is terminated --- .../java/com/firebolt/FireboltDriver.java | 33 ++++++++++++++++++- .../jdbc/connection/FireboltConnection.java | 20 +++++++---- .../FireboltConnectionServiceSecret.java | 10 +++--- .../FireboltConnectionUserPassword.java | 10 +++--- .../client/query/StatementClientImplTest.java | 3 +- .../FireboltConnectionServiceSecretTest.java | 4 +-- .../connection/FireboltConnectionTest.java | 3 ++ .../FireboltConnectionUserPasswordTest.java | 2 +- 8 files changed, 65 insertions(+), 20 deletions(-) diff --git a/src/main/java/com/firebolt/FireboltDriver.java b/src/main/java/com/firebolt/FireboltDriver.java index d8537117d..78f18440b 100644 --- a/src/main/java/com/firebolt/FireboltDriver.java +++ b/src/main/java/com/firebolt/FireboltDriver.java @@ -9,14 +9,20 @@ import java.sql.Driver; import java.sql.DriverPropertyInfo; import java.sql.SQLException; +import java.util.LinkedList; +import java.util.List; import java.util.Properties; +import java.util.logging.Level; import java.util.logging.Logger; +import static java.lang.String.format; + public class FireboltDriver implements Driver { public static final String JDBC_FIREBOLT = "jdbc:firebolt:"; private static final Logger rootLog; private static final Logger log; + private final List connections = new LinkedList<>(); static { try { @@ -29,9 +35,18 @@ public class FireboltDriver implements Driver { } } + public FireboltDriver() { + Runtime.getRuntime().addShutdownHook(new Thread(this::closeAllConnections)); + } + @Override public Connection connect(String url, Properties connectionSettings) throws SQLException { - return acceptsURL(url) ? FireboltConnection.create(url, connectionSettings) : null; + if (!acceptsURL(url)) { + return null; + } + Connection connection = FireboltConnection.create(url, connectionSettings, this); + connections.add(connection); + return connection; } @Override @@ -63,4 +78,20 @@ public boolean jdbcCompliant() { public Logger getParentLogger() { return rootLog; } + + public void removeClosedConnection(Connection connection) { + connections.remove(connection); + } + + private void closeAllConnections() { + for (Connection connection : connections) { + try { + if (!connection.isClosed()) { + connection.close(); + } + } catch (SQLException e) { + log.log(Level.WARNING, format("Cannot close connection on process shutting down %s", connection), e); + } + } + } } diff --git a/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java b/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java index e010c7284..296272658 100644 --- a/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java +++ b/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java @@ -1,5 +1,6 @@ package com.firebolt.jdbc.connection; +import com.firebolt.FireboltDriver; import com.firebolt.jdbc.JdbcBase; import com.firebolt.jdbc.annotation.ExcludeFromJacocoGeneratedReport; import com.firebolt.jdbc.annotation.NotImplemented; @@ -80,6 +81,7 @@ public abstract class FireboltConnection extends JdbcBase implements Connection, private final String protocolVersion; protected int infraVersion = 1; private DatabaseMetaData databaseMetaData; + private final FireboltDriver driver; //Properties that are used at the beginning of the connection for authentication protected final FireboltProperties loginProperties; @@ -89,7 +91,8 @@ protected FireboltConnection(@NonNull String url, Properties connectionSettings, FireboltAuthenticationService fireboltAuthenticationService, FireboltStatementService fireboltStatementService, - String protocolVersion) { + String protocolVersion, + FireboltDriver driver) { this.loginProperties = extractFireboltProperties(url, connectionSettings); this.fireboltAuthenticationService = fireboltAuthenticationService; @@ -100,11 +103,12 @@ protected FireboltConnection(@NonNull String url, this.connectionTimeout = loginProperties.getConnectionTimeoutMillis(); this.networkTimeout = loginProperties.getSocketTimeoutMillis(); this.protocolVersion = protocolVersion; + this.driver = driver; } // This code duplication between constructors is done because of back reference: dependent services require reference to current instance of FireboltConnection that prevents using constructor chaining or factory method. @ExcludeFromJacocoGeneratedReport - protected FireboltConnection(@NonNull String url, Properties connectionSettings, String protocolVersion) throws SQLException { + protected FireboltConnection(@NonNull String url, Properties connectionSettings, String protocolVersion, FireboltDriver driver) throws SQLException { this.loginProperties = extractFireboltProperties(url, connectionSettings); OkHttpClient httpClient = getHttpClient(loginProperties); @@ -116,18 +120,19 @@ protected FireboltConnection(@NonNull String url, Properties connectionSettings, this.connectionTimeout = loginProperties.getConnectionTimeoutMillis(); this.networkTimeout = loginProperties.getSocketTimeoutMillis(); this.protocolVersion = protocolVersion; + this.driver = driver; } protected abstract FireboltAuthenticationClient createFireboltAuthenticationClient(OkHttpClient httpClient); - public static FireboltConnection create(@NonNull String url, Properties connectionSettings) throws SQLException { - return createConnectionInstance(url, connectionSettings); + public static FireboltConnection create(@NonNull String url, Properties connectionSettings, FireboltDriver driver) throws SQLException { + return createConnectionInstance(url, connectionSettings, driver); } - private static FireboltConnection createConnectionInstance(@NonNull String url, Properties connectionSettings) throws SQLException { + private static FireboltConnection createConnectionInstance(@NonNull String url, Properties connectionSettings, FireboltDriver driver) throws SQLException { switch(getUrlVersion(url, connectionSettings)) { - case 1: return new FireboltConnectionUserPassword(url, connectionSettings); - case 2: return new FireboltConnectionServiceSecret(url, connectionSettings); + case 1: return new FireboltConnectionUserPassword(url, connectionSettings, driver); + case 2: return new FireboltConnectionServiceSecret(url, connectionSettings, driver); default: throw new IllegalArgumentException(format("Cannot distinguish version from url %s", url)); } } @@ -338,6 +343,7 @@ public void close() { statements.clear(); } databaseMetaData = null; + driver.removeClosedConnection(this); log.warning("Connection closed"); } diff --git a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java index e61092786..062a0f8f5 100644 --- a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java +++ b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java @@ -1,5 +1,6 @@ package com.firebolt.jdbc.connection; +import com.firebolt.FireboltDriver; import com.firebolt.jdbc.annotation.ExcludeFromJacocoGeneratedReport; import com.firebolt.jdbc.client.account.FireboltAccount; import com.firebolt.jdbc.client.account.FireboltAccountRetriever; @@ -43,8 +44,9 @@ public class FireboltConnectionServiceSecret extends FireboltConnection { FireboltGatewayUrlService fireboltGatewayUrlService, FireboltStatementService fireboltStatementService, FireboltEngineInformationSchemaService fireboltEngineService, - FireboltAccountIdService fireboltAccountIdService) throws SQLException { - super(url, connectionSettings, fireboltAuthenticationService, fireboltStatementService, PROTOCOL_VERSION); + FireboltAccountIdService fireboltAccountIdService, + FireboltDriver driver) throws SQLException { + super(url, connectionSettings, fireboltAuthenticationService, fireboltStatementService, PROTOCOL_VERSION, driver); this.fireboltGatewayUrlService = fireboltGatewayUrlService; this.fireboltAccountIdService = fireboltAccountIdService; this.fireboltEngineService = fireboltEngineService; @@ -52,8 +54,8 @@ public class FireboltConnectionServiceSecret extends FireboltConnection { } @ExcludeFromJacocoGeneratedReport - FireboltConnectionServiceSecret(@NonNull String url, Properties connectionSettings) throws SQLException { - super(url, connectionSettings, PROTOCOL_VERSION); + FireboltConnectionServiceSecret(@NonNull String url, Properties connectionSettings, FireboltDriver driver) throws SQLException { + super(url, connectionSettings, PROTOCOL_VERSION, driver); OkHttpClient httpClient = getHttpClient(loginProperties); this.fireboltGatewayUrlService = new FireboltGatewayUrlService(createFireboltAccountRetriever(httpClient,"engineUrl", GatewayUrlResponse.class)); this.fireboltAccountIdService = new FireboltAccountIdService(createFireboltAccountRetriever(httpClient,"resolve", FireboltAccount.class)); diff --git a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionUserPassword.java b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionUserPassword.java index e36845c1b..03739970d 100644 --- a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionUserPassword.java +++ b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionUserPassword.java @@ -1,5 +1,6 @@ package com.firebolt.jdbc.connection; +import com.firebolt.FireboltDriver; import com.firebolt.jdbc.annotation.ExcludeFromJacocoGeneratedReport; import com.firebolt.jdbc.client.account.FireboltAccountClient; import com.firebolt.jdbc.client.authentication.AuthenticationRequest; @@ -27,15 +28,16 @@ public class FireboltConnectionUserPassword extends FireboltConnection { Properties connectionSettings, FireboltAuthenticationService fireboltAuthenticationService, FireboltStatementService fireboltStatementService, - FireboltEngineInformationSchemaService fireboltEngineService) throws SQLException { - super(url, connectionSettings, fireboltAuthenticationService, fireboltStatementService, PROTOCOL_VERSION); + FireboltEngineInformationSchemaService fireboltEngineService, + FireboltDriver driver) throws SQLException { + super(url, connectionSettings, fireboltAuthenticationService, fireboltStatementService, PROTOCOL_VERSION, driver); this.fireboltEngineService = fireboltEngineService; connect(); } @ExcludeFromJacocoGeneratedReport - FireboltConnectionUserPassword(@NonNull String url, Properties connectionSettings) throws SQLException { - super(url, connectionSettings, PROTOCOL_VERSION); + FireboltConnectionUserPassword(@NonNull String url, Properties connectionSettings, FireboltDriver driver) throws SQLException { + super(url, connectionSettings, PROTOCOL_VERSION, driver); OkHttpClient httpClient = getHttpClient(loginProperties); this.fireboltEngineService = new FireboltEngineApiService(new FireboltAccountClient(httpClient, this, loginProperties.getUserDrivers(), loginProperties.getUserClients())); connect(); diff --git a/src/test/java/com/firebolt/jdbc/client/query/StatementClientImplTest.java b/src/test/java/com/firebolt/jdbc/client/query/StatementClientImplTest.java index eae589a79..519145a2c 100644 --- a/src/test/java/com/firebolt/jdbc/client/query/StatementClientImplTest.java +++ b/src/test/java/com/firebolt/jdbc/client/query/StatementClientImplTest.java @@ -1,5 +1,6 @@ package com.firebolt.jdbc.client.query; +import com.firebolt.FireboltDriver; import com.firebolt.jdbc.client.authentication.FireboltAuthenticationClient; import com.firebolt.jdbc.connection.FireboltConnection; import com.firebolt.jdbc.connection.FireboltConnectionTokens; @@ -349,7 +350,7 @@ private FireboltConnection use(int mockedInfraVersion, Properties props, String Call useCall = getMockedCallWithResponse(200, "", responseHeaders); Call select1Call = getMockedCallWithResponse(200, ""); when(okHttpClient.newCall(any())).thenReturn(useCall, select1Call); - FireboltConnection connection = new FireboltConnection("url", props, "0") { + FireboltConnection connection = new FireboltConnection("url", props, "0", mock(FireboltDriver.class)) { { this.infraVersion = mockedInfraVersion; try { diff --git a/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretTest.java b/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretTest.java index 4c02fc295..908a0a989 100644 --- a/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretTest.java +++ b/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretTest.java @@ -100,7 +100,7 @@ void checkSystemEngineEndpoint(String gatewayUrl, String expectedHost, String ex @SuppressWarnings("unchecked") FireboltAccountRetriever fireboltGatewayUrlClient = mock(FireboltAccountRetriever.class); when(fireboltGatewayUrlClient.retrieve(any(), any())).thenReturn(new GatewayUrlResponse(gatewayUrl)); FireboltGatewayUrlService gatewayUrlService = new FireboltGatewayUrlService(fireboltGatewayUrlClient); - FireboltConnection connection = new FireboltConnectionServiceSecret(SYSTEM_ENGINE_URL, connectionProperties, fireboltAuthenticationService, gatewayUrlService, fireboltStatementService, fireboltEngineService, fireboltAccountIdService); + FireboltConnection connection = new FireboltConnectionServiceSecret(SYSTEM_ENGINE_URL, connectionProperties, fireboltAuthenticationService, gatewayUrlService, fireboltStatementService, fireboltEngineService, fireboltAccountIdService, driver); FireboltProperties sessionProperties = connection.getSessionProperties(); assertEquals(expectedHost, sessionProperties.getHost()); assertEquals(expectedProps == null ? Map.of() : Arrays.stream(expectedProps.split(";")).map(kv -> kv.split("=")).collect(toMap(kv -> kv[0], kv -> kv[1])), sessionProperties.getAdditionalProperties()); @@ -113,6 +113,6 @@ void shouldNotFetchTokenNorEngineHostForLocalFirebolt() throws SQLException { } protected FireboltConnection createConnection(String url, Properties props) throws SQLException { - return new FireboltConnectionServiceSecret(url, props, fireboltAuthenticationService, fireboltGatewayUrlService, fireboltStatementService, fireboltEngineService, fireboltAccountIdService); + return new FireboltConnectionServiceSecret(url, props, fireboltAuthenticationService, fireboltGatewayUrlService, fireboltStatementService, fireboltEngineService, fireboltAccountIdService, driver); } } diff --git a/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionTest.java b/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionTest.java index c77710523..e842685b6 100644 --- a/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionTest.java +++ b/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionTest.java @@ -1,5 +1,6 @@ package com.firebolt.jdbc.connection; +import com.firebolt.FireboltDriver; import com.firebolt.jdbc.CheckedBiFunction; import com.firebolt.jdbc.CheckedFunction; import com.firebolt.jdbc.client.account.FireboltAccount; @@ -104,6 +105,8 @@ abstract class FireboltConnectionTest { protected FireboltAccountIdService fireboltAccountIdService; protected Properties connectionProperties = new Properties(); private static Connection connection; + @Mock + protected FireboltDriver driver; private final String URL; diff --git a/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionUserPasswordTest.java b/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionUserPasswordTest.java index 927fcdf34..3d6593e0d 100644 --- a/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionUserPasswordTest.java +++ b/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionUserPasswordTest.java @@ -73,6 +73,6 @@ void getMetadata(String engine) throws SQLException { } protected FireboltConnection createConnection(String url, Properties props) throws SQLException { - return new FireboltConnectionUserPassword(url, props, fireboltAuthenticationService, fireboltStatementService, fireboltEngineService); + return new FireboltConnectionUserPassword(url, props, fireboltAuthenticationService, fireboltStatementService, fireboltEngineService, driver); } } From c4170b58ff62593861f7de06b057ec091e60d810 Mon Sep 17 00:00:00 2001 From: alexradzin Date: Mon, 17 Jun 2024 16:37:02 +0300 Subject: [PATCH 2/3] fixed warnings --- .../jdbc/connection/FireboltConnectionServiceSecret.java | 1 + .../firebolt/jdbc/client/query/StatementClientImplTest.java | 6 +++--- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java index 062a0f8f5..58395eb9c 100644 --- a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java +++ b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java @@ -38,6 +38,7 @@ public class FireboltConnectionServiceSecret extends FireboltConnection { private final FireboltAccountIdService fireboltAccountIdService; private FireboltEngineService fireboltEngineService; // depends on infra version and is discovered during authentication + @SuppressWarnings("java:S107") // the price of the immutability FireboltConnectionServiceSecret(@NonNull String url, Properties connectionSettings, FireboltAuthenticationService fireboltAuthenticationService, diff --git a/src/test/java/com/firebolt/jdbc/client/query/StatementClientImplTest.java b/src/test/java/com/firebolt/jdbc/client/query/StatementClientImplTest.java index 519145a2c..a68e177d2 100644 --- a/src/test/java/com/firebolt/jdbc/client/query/StatementClientImplTest.java +++ b/src/test/java/com/firebolt/jdbc/client/query/StatementClientImplTest.java @@ -350,7 +350,7 @@ private FireboltConnection use(int mockedInfraVersion, Properties props, String Call useCall = getMockedCallWithResponse(200, "", responseHeaders); Call select1Call = getMockedCallWithResponse(200, ""); when(okHttpClient.newCall(any())).thenReturn(useCall, select1Call); - FireboltConnection connection = new FireboltConnection("url", props, "0", mock(FireboltDriver.class)) { + FireboltConnection conn = new FireboltConnection("url", props, "0", mock(FireboltDriver.class)) { { this.infraVersion = mockedInfraVersion; try { @@ -386,8 +386,8 @@ protected void assertDatabaseExisting(String database) { } }; - connection.createStatement().executeUpdate(useCommand); - return connection; + conn.createStatement().executeUpdate(useCommand); + return conn; } @ParameterizedTest From 6573a6fea6ce431951bd9989d6b5a45f9581ae31 Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Mon, 29 Jul 2024 11:11:12 +0300 Subject: [PATCH 3/3] fix build --- src/main/java/com/firebolt/FireboltDriver.java | 4 +--- .../java/com/firebolt/jdbc/connection/FireboltConnection.java | 2 +- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/src/main/java/com/firebolt/FireboltDriver.java b/src/main/java/com/firebolt/FireboltDriver.java index 1cf72e548..f2c9ec413 100644 --- a/src/main/java/com/firebolt/FireboltDriver.java +++ b/src/main/java/com/firebolt/FireboltDriver.java @@ -23,8 +23,6 @@ public class FireboltDriver implements Driver { public static final String JDBC_FIREBOLT = "jdbc:firebolt:"; - private static final Logger rootLog; - private static final Logger log; private final List connections = new LinkedList<>(); static { @@ -91,7 +89,7 @@ private void closeAllConnections() { connection.close(); } } catch (SQLException e) { - log.log(Level.WARNING, format("Cannot close connection on process shutting down %s", connection), e); + log.warn(format("Cannot close connection on process shutting down %s", connection), e); } } } diff --git a/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java b/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java index 34b00f123..993272057 100644 --- a/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java +++ b/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java @@ -343,7 +343,7 @@ public void close() { } databaseMetaData = null; driver.removeClosedConnection(this); - log.warning("Connection closed"); + log.warn("Connection closed"); } protected FireboltProperties extractFireboltProperties(String jdbcUri, Properties connectionProperties) {