add integration test

firebolt-db · Oct 29, 2024 · f105a2a · f105a2a
1 parent 8d4ff68
commit f105a2a
Show file tree

Hide file tree

Showing 4 changed files with 61 additions and 11 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,37 @@ void successfulConnect(boolean useDatabase, boolean useEngine) throws SQLExcepti
         }
     }
 
+    @Test
+    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;
@@ -61,7 +62,6 @@
 import java.util.function.Supplier;
 import java.util.regex.Pattern;
 
-import static com.firebolt.jdbc.connection.FireboltConnectionUserPassword.SYSTEM_ENGINE_NAME;
 import static com.firebolt.jdbc.connection.settings.FireboltSessionProperty.getNonDeprecatedProperties;
 import static java.lang.String.format;
 import static java.sql.ResultSet.CLOSE_CURSORS_AT_COMMIT;
@@ -428,23 +428,24 @@ public boolean isValid(int timeout) throws SQLException {
 			return false;
 		}
 		try {
-			validateConnection(getSessionProperties(), true, true, getSessionProperties().isValidateOnSystemEngine());
+			validateConnection(getSessionProperties(), true, true);
 			return true;
 		} catch (Exception e) {
 			return false;
 		}
 	}
 
-	private void validateConnection(FireboltProperties fireboltProperties, boolean ignoreToManyRequestsError, boolean isInternalRequest, boolean forceSystemEngine)
+	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");
 		}
-		if (forceSystemEngine) {
-			propertiesCopy.addProperty("engine", SYSTEM_ENGINE_NAME);
+		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.
@@ -487,7 +488,7 @@ private synchronized void changeProperty(Consumer<FireboltProperties> properties
 		try {
 			FireboltProperties tmpProperties = FireboltProperties.copy(sessionProperties);
 			propertiesEditor.accept(tmpProperties);
-			validateConnection(tmpProperties, false, false, false);
+			validateConnection(tmpProperties, false, false);
 			propertiesEditor.accept(sessionProperties);
 		} catch (FireboltException e) {
 			throw e;