added test that validates that system engine cannot execute queries o…

…n custom DB
firebolt-db · Nov 16, 2023 · 8872eaf · 8872eaf
1 parent 418ec79
commit 8872eaf
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 4 deletions.
diff --git a/src/integrationTest/java/integration/tests/SystemEngineTest.java b/src/integrationTest/java/integration/tests/SystemEngineTest.java
@@ -1,6 +1,7 @@
 package integration.tests;
 
 import com.firebolt.jdbc.exception.FireboltException;
+import integration.ConnectionInfo;
 import integration.IntegrationTest;
 import lombok.CustomLog;
 import org.junit.jupiter.api.AfterAll;
@@ -11,6 +12,7 @@
 import org.junit.jupiter.api.TestMethodOrder;
 
 import java.sql.Connection;
+import java.sql.DriverManager;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
@@ -51,7 +53,7 @@ void afterAll() {
 
 	@Test
 	void shouldSelect1() throws SQLException {
-		try (Connection connection = this.createConnection(getSystemEngineName());
+		try (Connection connection = createConnection(getSystemEngineName());
 			 ResultSet rs = connection.createStatement().executeQuery("SELECT 1")) {
 			assertTrue(rs.next());
 			assertEquals(1, rs.getInt(1));
@@ -61,7 +63,7 @@ void shouldSelect1() throws SQLException {
 
 	@Test
 	void ddl() throws SQLException {
-		try (Connection connection = this.createConnection(getSystemEngineName())) {
+		try (Connection connection = createConnection(getSystemEngineName())) {
 			assertEquals(
 					"Cannot execute a DDL query on the system engine.",
 					assertThrows(
@@ -71,6 +73,33 @@ void ddl() throws SQLException {
 		}
 	}
 
+	@Test
+	void shouldFailToSelectFromCustomDbUsingSystemEngine() throws SQLException {
+		ConnectionInfo current = integration.ConnectionInfo.getInstance();
+		String systemEngineJdbcUrl = new ConnectionInfo(current.getPrincipal(), current.getSecret(),
+				current.getEnv(), current.getDatabase(), current.getAccount(), getSystemEngineName(), current.getApi()).toJdbcUrl();
+		String customEngineJdbcUrl = current.toJdbcUrl();
+		String principal = current.getPrincipal();
+		String secret = current.getSecret();
+
+		try (Connection systemConnection = DriverManager.getConnection(systemEngineJdbcUrl, principal, secret);
+			 Connection customConnection = DriverManager.getConnection(customEngineJdbcUrl, principal, secret)) {
+
+			try {
+				customConnection.createStatement().executeUpdate("CREATE DIMENSION TABLE dummy(id INT)");
+				try (ResultSet rs = customConnection.createStatement().executeQuery("select count(*) from dummy")) {
+					assertTrue(rs.next());
+					assertEquals(0, rs.getInt(1));
+					assertFalse(rs.next());
+				}
+				FireboltException e = assertThrows(FireboltException.class, () -> systemConnection.createStatement().executeQuery("select count(*) from dummy"));
+				assertEquals("Queries against table dummy require a user engine", e.getErrorMessageFromServer().replaceAll("\r?\n", ""));
+			} finally {
+				customConnection.createStatement().executeUpdate("DROP TABLE dummy");
+			}
+		}
+	}
+
 	@Test
 	@Tag("slow")
 	void shouldExecuteEngineManagementQueries() throws SQLException {

diff --git a/src/main/java/com/firebolt/jdbc/client/query/StatementClientImpl.java b/src/main/java/com/firebolt/jdbc/client/query/StatementClientImpl.java
@@ -221,12 +221,12 @@ private Map<String, String> getAllParameters(FireboltProperties fireboltProperti
 		} else {
 			params.put(FireboltQueryParameterKey.QUERY_ID.getKey(), statementInfoWrapper.getId());
 			params.put(FireboltQueryParameterKey.COMPRESS.getKey(), fireboltProperties.isCompress() ? "1" : "0");
-			params.put(FireboltQueryParameterKey.DATABASE.getKey(), fireboltProperties.getDatabase());
 
 			if (queryTimeout > 0) {
 				params.put("max_execution_time", String.valueOf(queryTimeout));
 			}
 		}
+		params.put(FireboltQueryParameterKey.DATABASE.getKey(), fireboltProperties.getDatabase());
 
 		return params;
 	}

diff --git a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java
@@ -98,7 +98,8 @@ private FireboltProperties getSessionPropertiesForSystemEngine(String accessToke
                 .additionalProperties(Map.of())
                 .compress(false)
                 .accountId(accountId)
-                .host(UrlUtil.createUrl(systemEngineEndpoint).getHost()).database(null).build();
+                .host(UrlUtil.createUrl(systemEngineEndpoint).getHost())
+                .build();
     }
 
     @Override