FIR-33849 close all open connections when process is terminated

src/main/java/com/firebolt/FireboltDriver.java

@@ -10,14 +10,20 @@
 import java.sql.Driver;
 import java.sql.DriverPropertyInfo;
 import java.sql.SQLException;
+import java.util.LinkedList;
+import java.util.List;
 import java.sql.SQLFeatureNotSupportedException;
 import java.util.Properties;
+import java.util.logging.Level;
 import java.util.logging.Logger;
 
+import static java.lang.String.format;
+
 @CustomLog
 public class FireboltDriver implements Driver {
 
 	public static final String JDBC_FIREBOLT = "jdbc:firebolt:";
+	private final List<Connection> connections = new LinkedList<>();
 
 	static {
 		try {
@@ -28,9 +34,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
@@ -62,4 +77,20 @@ public boolean jdbcCompliant() {
 	public Logger getParentLogger() throws SQLFeatureNotSupportedException {
 		throw new FireboltSQLFeatureNotSupportedException();
 	}
+
+	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.warn(format("Cannot close connection on process shutting down %s", connection), e);
+			}
+		}
+	}
 }

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;
@@ -79,6 +80,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;
@@ -88,7 +90,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;
@@ -99,11 +102,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);
 
@@ -115,18 +119,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));
 		}
 	}
@@ -337,7 +342,8 @@ public void close() {
 			statements.clear();
 		}
 		databaseMetaData = null;
-		log.debug("Connection closed");
+		driver.removeClosedConnection(this);
+		log.warn("Connection closed");
 	}
 
 	protected FireboltProperties extractFireboltProperties(String jdbcUri, Properties connectionProperties) {

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;
@@ -37,23 +38,25 @@ 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,
                                     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;
         connect();
     }
 
     @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));

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

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 conn = new FireboltConnection("url", props, "0", mock(FireboltDriver.class)) {
 			{
 				this.infraVersion = mockedInfraVersion;
 				try {
@@ -385,8 +386,8 @@ protected void assertDatabaseExisting(String database) {
 
 			}
 		};
-		connection.createStatement().executeUpdate(useCommand);
-		return connection;
+		conn.createStatement().executeUpdate(useCommand);
+		return conn;
 	}
 
 	@ParameterizedTest

src/test/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretTest.java

@@ -100,7 +100,7 @@ void checkSystemEngineEndpoint(String gatewayUrl, String expectedHost, String ex
         @SuppressWarnings("unchecked") FireboltAccountRetriever<GatewayUrlResponse> 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);
     }
 }

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;
 

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