attempt to fix system engine tests for old auth - does not work

firebolt-db · Nov 14, 2023 · 85d68da · 85d68da
1 parent 57f00ad
commit 85d68da
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 21 deletions.
diff --git a/src/integrationTest/java/integration/tests/StatementTest.java b/src/integrationTest/java/integration/tests/StatementTest.java
@@ -31,21 +31,23 @@
 @CustomLog
 class StatementTest extends IntegrationTest {
 
-	@BeforeEach
-	void beforeEach() {
-		executeStatementFromFile("/statements/statement/ddl.sql");
-	}
-
-	@AfterEach
-	void afterEach() {
-		executeStatementFromFile("/statements/statement/cleanup.sql");
-	}
+//	@BeforeEach
+//	void beforeEach() {
+//		executeStatementFromFile("/statements/statement/ddl.sql");
+//	}
+//
+//	@AfterEach
+//	void afterEach() {
+//		executeStatementFromFile("/statements/statement/cleanup.sql");
+//	}
 
 	@Test
 	void shouldSelect1() throws SQLException {
-		try (Connection connection = this.createConnection(); Statement statement = connection.createStatement()) {
-			statement.executeQuery("SELECT 1;");
-			assertNotNull(statement.executeQuery("SELECT 1;"));
+		try (Connection connection = this.createConnection();
+			 ResultSet rs = connection.createStatement().executeQuery("SELECT 1")) {
+			assertTrue(rs.next());
+			assertEquals(1, rs.getInt(1));
+			assertFalse(rs.next());
 		}
 	}
 

diff --git a/src/integrationTest/java/integration/tests/SystemEngineTest.java b/src/integrationTest/java/integration/tests/SystemEngineTest.java
@@ -1,15 +1,26 @@
 package integration.tests;
 
+import com.firebolt.jdbc.exception.FireboltException;
+import integration.IntegrationTest;
+import lombok.CustomLog;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.MethodOrderer;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestMethodOrder;
+
 import java.sql.Connection;
+import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.Arrays;
 import java.util.List;
 
-import org.junit.jupiter.api.*;
-
-import integration.IntegrationTest;
-import lombok.CustomLog;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 @CustomLog
 @TestMethodOrder(MethodOrderer.OrderAnnotation.class)
@@ -37,10 +48,32 @@ void afterAll() {
 		}
 	}
 
+	@Test
+	void shouldSelect1() throws SQLException {
+		try (Connection connection = this.createConnection(getSystemEngineName());
+			 ResultSet rs = connection.createStatement().executeQuery("SELECT 1")) {
+			assertTrue(rs.next());
+			assertEquals(1, rs.getInt(1));
+			assertFalse(rs.next());
+		}
+	}
+
+	@Test
+	void ddl() throws SQLException {
+		try (Connection connection = this.createConnection(getSystemEngineName())) {
+			assertEquals(
+					"Cannot execute a DDL query on the system engine.",
+					assertThrows(
+							FireboltException.class,
+							() -> connection.createStatement().executeUpdate("CREATE DIMENSION TABLE dummy(id INT)"))
+							.getErrorMessageFromServer().replaceAll("\r?\n", ""));
+		}
+	}
+
 	@Test
 	@Tag("slow")
 	void shouldExecuteEngineManagementQueries() throws SQLException {
-		try (Connection connection = this.createConnection(null)) {
+		try (Connection connection = this.createConnection(getSystemEngineName())) {
 			List<String> queries = Arrays.asList(String.format("CREATE DATABASE IF NOT EXISTS %s", DATABASE_NAME),
 					String.format("CREATE ENGINE %s", ENGINE_NAME),
 					String.format("ATTACH ENGINE %s TO %s;", ENGINE_NAME, DATABASE_NAME),
@@ -55,4 +88,8 @@ void shouldExecuteEngineManagementQueries() throws SQLException {
 			}
 		}
 	}
+
+	private String getSystemEngineName() {
+		return System.getProperty("api") == null ? null : "system";
+	}
 }
diff --git a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionUserPasswordAuthentication.java b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionUserPasswordAuthentication.java
@@ -21,6 +21,8 @@
 import java.sql.SQLException;
 import java.util.Properties;
 
+import static java.lang.String.format;
+
 public class FireboltConnectionUserPasswordAuthentication extends FireboltConnection {
     private final FireboltEngineService fireboltEngineService;
 
@@ -45,10 +47,14 @@ public class FireboltConnectionUserPasswordAuthentication extends FireboltConnec
 
     @Override
     protected void authenticate() throws SQLException {
-        String accessToken = getAccessToken(loginProperties).orElse(StringUtils.EMPTY);
-        FireboltProperties propertiesWithAccessToken = loginProperties.toBuilder().accessToken(accessToken).build();
-        Engine engine = fireboltEngineService.getEngine(propertiesWithAccessToken);
-        this.sessionProperties = loginProperties.toBuilder().host(engine.getEndpoint()).engine(engine.getName()).build();
+        if (loginProperties.isSystemEngine()) {
+            this.sessionProperties = loginProperties;
+        } else {
+            String accessToken = getAccessToken(loginProperties).orElse(StringUtils.EMPTY);
+            FireboltProperties propertiesWithAccessToken = loginProperties.toBuilder().accessToken(accessToken).build();
+            Engine engine = fireboltEngineService.getEngine(propertiesWithAccessToken);
+            this.sessionProperties = loginProperties.toBuilder().host(engine.getEndpoint()).engine(engine.getName()).build();
+        }
     }
 
     @Override