Fir 37710 jdbc validate connection against system engine (#473)

firebolt-db · Oct 30, 2024 · 8f8a82d · 8f8a82d
1 parent 7c50295
commit 8f8a82d
Showing 10 changed files with 113 additions and 27 deletions.
diff --git a/src/integrationTest/java/integration/ConnectionInfo.java b/src/integrationTest/java/integration/ConnectionInfo.java
@@ -1,5 +1,7 @@
 package integration;
 
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.function.Supplier;
@@ -20,6 +22,7 @@ public class ConnectionInfo {
 	private final String account;
 	private final String engine;
 	private final String api;
+	private final Map<String, String> extra;
 	private final Supplier<String> jdbcUrlSupplier;
 
 	private ConnectionInfo() {
@@ -35,13 +38,18 @@ private ConnectionInfo() {
 	}
 
 	public ConnectionInfo(String principal, String secret, String env, String database, String account, String engine, String api) {
+		this(principal, secret, env, database, account, engine, api, new HashMap<>());
+	}
+
+	public ConnectionInfo(String principal, String secret, String env, String database, String account, String engine, String api, Map<String, String> extra) {
 		this.principal = principal;
 		this.secret = secret;
 		this.env = env;
 		this.database = database;
 		this.account = account;
 		this.engine = engine;
 		this.api = api;
+		this.extra = extra;
 		jdbcUrlSupplier = api == null ? this::toJdbcUrl2 : this::toJdbcUrl1;
 	}
 
@@ -97,7 +105,10 @@ private String toJdbcUrl1() {
 	}
 
 	private String toJdbcUrl2() {
-		String params = Stream.of(param("env", env), param("engine", engine), param("account", account)).filter(Objects::nonNull).collect(joining("&"));
+		String params = Stream.concat(
+				Stream.of(param("env", env), param("engine", engine), param("account", account)),
+				extra.entrySet().stream().map(e -> param(e.getKey(), e.getValue()))
+		).filter(Objects::nonNull).collect(joining("&"));
 		if (!params.isEmpty()) {
 			params = "?" + params;
 		}

diff --git a/src/integrationTest/java/integration/IntegrationTest.java b/src/integrationTest/java/integration/IntegrationTest.java
@@ -13,6 +13,8 @@
 import java.sql.DriverManager;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.HashMap;
+import java.util.Map;
 
 import static com.firebolt.jdbc.connection.FireboltConnectionUserPassword.SYSTEM_ENGINE_NAME;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -37,9 +39,13 @@ protected Connection createConnection() throws SQLException {
 	}
 
 	protected Connection createConnection(String engine) throws SQLException {
+		return createConnection(engine, new HashMap<>());
+	}
+
+	protected Connection createConnection(String engine, Map<String, String> extra) throws SQLException {
 		ConnectionInfo current = integration.ConnectionInfo.getInstance();
 		ConnectionInfo updated = new ConnectionInfo(current.getPrincipal(), current.getSecret(),
-				current.getEnv(), current.getDatabase(), current.getAccount(), engine, current.getApi());
+				current.getEnv(), current.getDatabase(), current.getAccount(), engine, current.getApi(), extra);
 		return DriverManager.getConnection(updated.toJdbcUrl(),
 				integration.ConnectionInfo.getInstance().getPrincipal(),
 				integration.ConnectionInfo.getInstance().getSecret());

diff --git a/src/integrationTest/java/integration/tests/ConnectionTest.java b/src/integrationTest/java/integration/tests/ConnectionTest.java
@@ -19,6 +19,7 @@
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.Map;
 
 import static java.lang.String.format;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -139,6 +140,38 @@ void successfulConnect(boolean useDatabase, boolean useEngine) throws SQLExcepti
         }
     }
 
+    @Test
+    @Tag("v2")
+    void validatesOnSystemEngineIfParameterProvided() throws SQLException {
+        try (Connection systemConnection = createConnection(null)) {
+            String engineName = integration.ConnectionInfo.getInstance().getEngine() + "_validate_test";
+            try (Statement systemStatement = systemConnection.createStatement()) {
+                systemStatement.executeUpdate(format("CREATE ENGINE %s WITH INITIALLY_STOPPED=true", engineName));
+            }
+            try (Connection connection = createConnection(engineName, Map.of("validate_on_system_engine", "true"))) {
+                try (Statement systemStatement = systemConnection.createStatement()) {
+                    ResultSet rs = systemStatement.executeQuery(
+                            format("SELECT status FROM information_schema.engines WHERE engine_name='%s'", engineName));
+                    assertTrue(rs.next());
+                    assertEquals("STOPPED", rs.getString(1));
+                }
+                assertTrue(connection.isValid(500));
+                // After validation the engine should still be stopped
+                try (Statement systemStatement = systemConnection.createStatement()) {
+                    ResultSet rs = systemStatement.executeQuery(
+                            format("SELECT status FROM information_schema.engines WHERE engine_name='%s'", engineName));
+                    assertTrue(rs.next());
+                    assertEquals("STOPPED", rs.getString(1));
+                }
+            } finally {
+                try (Statement systemStatement = systemConnection.createStatement()) {
+                    systemStatement.executeUpdate(format("DROP ENGINE %s", engineName));
+                }
+            }
+        }
+
+    }
+
     void unsuccessfulConnect(boolean useDatabase, boolean useEngine) throws SQLException {
         ConnectionInfo params = integration.ConnectionInfo.getInstance();
         String url = getJdbcUrl(params, useDatabase, useEngine);

diff --git a/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java b/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java
@@ -51,6 +51,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.IdentityHashMap;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -427,9 +428,7 @@ public boolean isValid(int timeout) throws SQLException {
 			return false;
 		}
 		try {
-			if (!loginProperties.isSystemEngine()) {
-				validateConnection(getSessionProperties(), true, true);
-			}
+			validateConnection(getSessionProperties(), true, true);
 			return true;
 		} catch (Exception e) {
 			return false;
@@ -438,11 +437,15 @@ public boolean isValid(int timeout) throws SQLException {
 
 	private void validateConnection(FireboltProperties fireboltProperties, boolean ignoreToManyRequestsError, boolean isInternalRequest)
 			throws SQLException {
-		FireboltProperties propertiesCopy = FireboltProperties.copy(fireboltProperties);
+		HashMap<String, String> runtimeProperties = new HashMap<>(fireboltProperties.getRuntimeAdditionalProperties());
 		if (isInternalRequest) {
-			propertiesCopy.addProperty("auto_start_stop_control", "ignore");
+			runtimeProperties.put("auto_start_stop_control", "ignore");
+		}
+		var propertiesBuilder = fireboltProperties.toBuilder().runtimeAdditionalProperties(runtimeProperties);
+		if (getSessionProperties().isValidateOnSystemEngine()) {
+			propertiesBuilder.compress(false).engine(null).systemEngine(true);
 		}
-		try (Statement s = createStatement(propertiesCopy)) {
+		try (Statement s = createStatement(propertiesBuilder.build())) {
 			s.execute("SELECT 1");
 		} catch (Exception e) {
 			// A connection is not invalid when too many requests are being sent.

diff --git a/src/main/java/com/firebolt/jdbc/connection/settings/FireboltProperties.java b/src/main/java/com/firebolt/jdbc/connection/settings/FireboltProperties.java
@@ -75,6 +75,7 @@ public class FireboltProperties {
 	private final String userDrivers;
 	private final String userClients;
 	private final String accessToken;
+	private final boolean validateOnSystemEngine;
 	@Builder.Default
 	private Map<String, String> initialAdditionalProperties = new HashMap<>();
 	@Builder.Default
@@ -110,6 +111,7 @@ public FireboltProperties(Properties properties) {
 		String configuredEnvironment = getSetting(properties, FireboltSessionProperty.ENVIRONMENT);
 		userDrivers = getSetting(properties, FireboltSessionProperty.USER_DRIVERS);
 		userClients = getSetting(properties, FireboltSessionProperty.USER_CLIENTS);
+		validateOnSystemEngine = getSetting(properties, FireboltSessionProperty.VALIDATE_ON_SYSTEM_ENGINE);
 
 		environment = getEnvironment(configuredEnvironment, properties);
 		host = getHost(configuredEnvironment, properties);

diff --git a/src/main/java/com/firebolt/jdbc/connection/settings/FireboltSessionProperty.java b/src/main/java/com/firebolt/jdbc/connection/settings/FireboltSessionProperty.java
@@ -67,6 +67,9 @@ public enum FireboltSessionProperty {
 	USER_CLIENTS("user_clients", null, String.class, "user clients", FireboltProperties::getUserClients),
 	ACCESS_TOKEN("access_token", null, String.class, "access token", p -> "***"),
 	ENVIRONMENT("environment", "app", String.class, "Firebolt environment", FireboltProperties::getEnvironment, "env"),
+	VALIDATE_ON_SYSTEM_ENGINE("validate_on_system_engine", false, Boolean.class,
+			"Whether to validate the connection on the system engine or not. By default validates on an engine currently connected.",
+			FireboltProperties::isValidateOnSystemEngine),
 	// We keep all the deprecated properties to ensure backward compatibility - but
 	// they do not have any effect.
 	@Deprecated

diff --git a/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretTest.java b/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretTest.java
@@ -39,15 +39,6 @@ public FireboltConnectionServiceSecretTest() {
         super("jdbc:firebolt:db?env=dev&engine=eng&account=dev");
     }
 
-    @Test
-    void shouldNotValidateConnectionWhenCallingIsValidWhenUsingSystemEngine() throws SQLException {
-        Properties propertiesWithSystemEngine = new Properties(connectionProperties);
-        try (FireboltConnection fireboltConnection = createConnection(SYSTEM_ENGINE_URL, propertiesWithSystemEngine)) {
-            fireboltConnection.isValid(500);
-            verifyNoInteractions(fireboltStatementService);
-        }
-    }
-
     @Test
     void shouldNotGetEngineUrlOrDefaultEngineUrlWhenUsingSystemEngine() throws SQLException {
         connectionProperties.put("database", "my_db");

diff --git a/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionTest.java b/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionTest.java
@@ -51,6 +51,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Stream;
 
+import static com.firebolt.jdbc.connection.FireboltConnectionUserPassword.SYSTEM_ENGINE_NAME;
 import static com.firebolt.jdbc.connection.settings.FireboltSessionProperty.ACCESS_TOKEN;
 import static com.firebolt.jdbc.connection.settings.FireboltSessionProperty.CLIENT_ID;
 import static com.firebolt.jdbc.connection.settings.FireboltSessionProperty.CLIENT_SECRET;
@@ -86,6 +87,7 @@
 @ExtendWith(MockitoExtension.class)
 abstract class FireboltConnectionTest {
 	private static final String LOCAL_URL = "jdbc:firebolt:local_dev_db?account=dev&ssl=false&max_query_size=10000000&mask_internal_errors=0&host=localhost";
+	private static final String SYSTEM_ENGINE_URL = "jdbc:firebolt:db?env=dev&account=dev";
 	private final FireboltConnectionTokens fireboltConnectionTokens = new FireboltConnectionTokens(null, 0);
 	@Captor
 	private ArgumentCaptor<FireboltProperties> propertiesArgumentCaptor;
@@ -340,6 +342,37 @@ void shouldValidateConnectionWhenCallingIsValid() throws SQLException {
 		}
 	}
 
+	@Test
+	void shouldValidateConnectionWhenCallingIsValidSystemEngine() throws SQLException {
+		when(fireboltStatementService.execute(any(), any(), any()))
+				.thenReturn(Optional.empty());
+		Properties propertiesWithSystemEngine = new Properties(connectionProperties);
+		try (FireboltConnection fireboltConnection = createConnection(SYSTEM_ENGINE_URL, propertiesWithSystemEngine)) {
+			fireboltConnection.isValid(500);
+			verify(fireboltStatementService).execute(queryInfoWrapperArgumentCaptor.capture(),
+					propertiesArgumentCaptor.capture(), any());
+			assertEquals(List.of("SELECT 1"), queryInfoWrapperArgumentCaptor.getAllValues().stream().map(StatementInfoWrapper::getSql).collect(toList()));
+			assertEquals(Map.of("auto_start_stop_control", "ignore"), propertiesArgumentCaptor.getValue().getAdditionalProperties());
+		}
+	}
+
+	@Test
+	void shouldForceSystemEngineWhenValidateOnSystemEngineIsSet() throws SQLException {
+		when(fireboltStatementService.execute(any(), any(), any()))
+				.thenReturn(Optional.empty());
+		Properties propertiesCopy = new Properties(connectionProperties);
+		propertiesCopy.put("validate_on_system_engine", "true");
+		try (FireboltConnection fireboltConnection = createConnection(URL, propertiesCopy)) {
+			fireboltConnection.isValid(500);
+			verify(fireboltStatementService).execute(queryInfoWrapperArgumentCaptor.capture(),
+					propertiesArgumentCaptor.capture(), any());
+			assertEquals(List.of("SELECT 1"), queryInfoWrapperArgumentCaptor.getAllValues().stream().map(StatementInfoWrapper::getSql).collect(toList()));
+			assertEquals(Map.of("auto_start_stop_control", "ignore"), propertiesArgumentCaptor.getValue().getAdditionalProperties());
+			assertEquals(null, propertiesArgumentCaptor.getValue().getEngine());
+			assertTrue(propertiesArgumentCaptor.getValue().isSystemEngine());
+		}
+	}
+
 	@Test
 	void shouldIgnore429WhenValidatingConnection() throws SQLException {
 		when(fireboltStatementService.execute(any(), any(), any()))
@@ -694,5 +727,16 @@ void getClientInfo() throws SQLException {
 		}
 	}
 
+	@Test
+	void shouldValidateOnUserEngineByDefault() throws SQLException {
+		try (FireboltConnection fireboltConnection = createConnection(URL, connectionProperties)) {
+			assertEquals("false", fireboltConnection.getClientInfo().get("validate_on_system_engine"));
+		}
+		connectionProperties.put("validate_on_system_engine", "true");
+		try (FireboltConnection fireboltConnection = createConnection(URL, connectionProperties)) {
+			assertEquals("true", fireboltConnection.getClientInfo().get("validate_on_system_engine"));
+		}
+	}
+
 	protected abstract FireboltConnection createConnection(String url, Properties props) throws SQLException;
 }
diff --git a/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionUserPasswordTest.java b/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionUserPasswordTest.java
@@ -30,15 +30,6 @@ public FireboltConnectionUserPasswordTest() {
         super("jdbc:firebolt://api.dev.firebolt.io/db");
     }
 
-    @Test
-    void shouldNotValidateConnectionWhenCallingIsValidWhenUsingSystemEngine() throws SQLException {
-        Properties propertiesWithSystemEngine = new Properties(connectionProperties);
-        try (FireboltConnection fireboltConnection = createConnection(SYSTEM_ENGINE_URL, propertiesWithSystemEngine)) {
-            fireboltConnection.isValid(500);
-            verifyNoInteractions(fireboltStatementService);
-        }
-    }
-
     @Test
     void shouldNotGetEngineUrlOrDefaultEngineUrlWhenUsingSystemEngine() throws SQLException {
         connectionProperties.put("database", "my_db");

diff --git a/src/test/java/com/firebolt/jdbc/connection/settings/FireboltPropertiesTest.java b/src/test/java/com/firebolt/jdbc/connection/settings/FireboltPropertiesTest.java
@@ -48,6 +48,7 @@ void shouldHaveAllTheSpecifiedCustomProperties() {
 		properties.put("path", "example");
 		properties.put("someCustomProperties", "custom_value");
 		properties.put("compress", "1");
+		properties.put("validate_on_system_engine", "true");
 
 		Map<String, String> customProperties = new HashMap<>();
 		customProperties.put("someCustomProperties", "custom_value");
@@ -57,7 +58,8 @@ void shouldHaveAllTheSpecifiedCustomProperties() {
 				.port(443).principal(null).secret(null).host("myDummyHost").ssl(true).systemEngine(false)
 				.initialAdditionalProperties(customProperties).keepAliveTimeoutMillis(300000)
 				.maxConnectionsTotal(300).maxRetries(3).socketTimeoutMillis(20).connectionTimeoutMillis(60000)
-				.tcpKeepInterval(30).tcpKeepIdle(60).tcpKeepCount(10).environment("app").build();
+				.tcpKeepInterval(30).tcpKeepIdle(60).tcpKeepCount(10).environment("app").validateOnSystemEngine(true)
+				.build();
 		assertEquals(expectedDefaultProperties, new FireboltProperties(properties));
 	}