Fix using the shared/global connection in IntegrationTestUtil (#494)

All the integration test use the static `Connection` object in `IntegrationTestUtil` to execute query. There are 2 issues
  - TestUtil should not have any static mutable member
  - The connection member will be overwritten by lots of the test and will result corrupted state
  
We should pass the connection object as the argument of the execute API in the util.

Fix stub mappings
---------

Co-authored-by: Jayant Singh <[email protected]>

Author: jackyhu-db

src/test/java/com/databricks/jdbc/integration/IntegrationTestUtil.java

@@ -24,8 +24,6 @@ public class IntegrationTestUtil {
   private static final boolean isFakeServiceTest =
       Boolean.parseBoolean(System.getProperty(IS_FAKE_SERVICE_TEST_PROP));
 
-  private static Connection JDBCConnection;
-
   /** Get the host of the embedded web server of fake service to be used in the tests. */
   public static String getFakeServiceHost() {
     // Target base URL of the fake service type
@@ -166,10 +164,6 @@ public static Connection getBenchfoodJDBCConnection() throws SQLException {
         getBenchfoodJDBCUrl(), getDatabricksUser(), getDatabricksBenchfoodToken());
   }
 
-  public static void resetJDBCConnection() {
-    JDBCConnection = null;
-  }
-
   public static String getJDBCUrl() {
     String template = "jdbc:databricks://%s/default;ssl=1;AuthMech=3;httpPath=%s";
     String host = getDatabricksHost();
@@ -213,89 +207,88 @@ public static String getDogfoodJDBCUrl() {
     return String.format(template, host, httpPath);
   }
 
-  public static boolean executeSQL(String sql) {
+  public static boolean executeSQL(Connection conn, String sql) {
     try {
-      if (JDBCConnection == null) JDBCConnection = getValidJDBCConnection();
-      JDBCConnection.createStatement().execute(sql);
+      conn.createStatement().execute(sql);
       return true;
     } catch (SQLException e) {
       System.out.println("Error executing SQL: " + e.getMessage());
       return false;
     }
   }
 
-  public static ResultSet executeQuery(String sql) {
+  public static ResultSet executeQuery(Connection conn, String sql) {
     try {
-      if (JDBCConnection == null) JDBCConnection = getValidJDBCConnection();
-      ResultSet rs = JDBCConnection.createStatement().executeQuery(sql);
+      ResultSet rs = conn.createStatement().executeQuery(sql);
       return rs;
     } catch (SQLException e) {
       System.out.println("Error executing SQL: " + e.getMessage());
       return null;
     }
   }
 
-  public static void setupDatabaseTable(String tableName) {
+  public static void setupDatabaseTable(Connection conn, String tableName) {
     String tableDeletionSQL = "DROP TABLE IF EXISTS " + getFullyQualifiedTableName(tableName);
-    executeSQL(tableDeletionSQL);
+    executeSQL(conn, tableDeletionSQL);
 
     String tableCreationSQL =
         "CREATE TABLE IF NOT EXISTS "
             + getFullyQualifiedTableName(tableName)
             + " (id INT PRIMARY KEY, col1 VARCHAR(255), col2 VARCHAR(255))";
 
-    executeSQL(tableCreationSQL);
+    executeSQL(conn, tableCreationSQL);
   }
 
-  public static void setupDatabaseTable(String tableName, String tableCreationSQL) {
+  public static void setupDatabaseTable(
+      Connection conn, String tableName, String tableCreationSQL) {
     String tableDeletionSQL = "DROP TABLE IF EXISTS " + getFullyQualifiedTableName(tableName);
 
-    executeSQL(tableDeletionSQL);
-    executeSQL(tableCreationSQL);
+    executeSQL(conn, tableDeletionSQL);
+    executeSQL(conn, tableCreationSQL);
   }
 
-  public static void deleteTable(String tableName) {
+  public static void deleteTable(Connection conn, String tableName) {
     String SQL = "DROP TABLE IF EXISTS " + getFullyQualifiedTableName(tableName);
-    executeSQL(SQL);
+    executeSQL(conn, SQL);
   }
 
   public static String getFullyQualifiedTableName(String tableName) {
     return getDatabricksCatalog() + "." + getDatabricksSchema() + "." + tableName;
   }
 
-  public static void insertTestDataForJoins(String table1Name, String table2Name) {
+  public static void insertTestDataForJoins(Connection conn, String table1Name, String table2Name) {
     // Insert data into the first table
     String insertTable1SQL1 =
         "INSERT INTO "
             + getFullyQualifiedTableName(table1Name)
             + " (id, col1, col2) VALUES (1, 'value1_table1', 'value2_table1')";
-    executeSQL(insertTable1SQL1);
+    executeSQL(conn, insertTable1SQL1);
 
     String insertTable1SQL2 =
         "INSERT INTO "
             + getFullyQualifiedTableName(table1Name)
             + " (id, col1, col2) VALUES (2, 'value3_table1', 'value4_table1')";
-    executeSQL(insertTable1SQL2);
+    executeSQL(conn, insertTable1SQL2);
 
     // Insert related data into the second table
     String insertTable2SQL1 =
         "INSERT INTO "
             + getFullyQualifiedTableName(table2Name)
             + " (id, col1, col2) VALUES (1, 'related_value1_table2', 'related_value2_table2')";
-    executeSQL(insertTable2SQL1);
+    executeSQL(conn, insertTable2SQL1);
 
     String insertTable2SQL2 =
         "INSERT INTO "
             + getFullyQualifiedTableName(table2Name)
             + " (id, col1, col2) VALUES (2, 'related_value3_table2', 'related_value4_table2')";
-    executeSQL(insertTable2SQL2);
+    executeSQL(conn, insertTable2SQL2);
   }
 
-  public static void insertTestData(String tableName) {
+  public static void insertTestData(Connection conn, String tableName) {
     String insertSQL =
         "INSERT INTO "
             + getFullyQualifiedTableName(tableName)
             + " (id, col1, col2) VALUES (1, 'value1', 'value2')";
-    executeSQL(insertSQL);
+    executeSQL(conn, insertSQL);
   }
 }

src/test/java/com/databricks/jdbc/integration/e2e/LargeQueryExecutionTests.java

@@ -3,17 +3,34 @@
 import static com.databricks.jdbc.integration.IntegrationTestUtil.*;
 import static org.junit.jupiter.api.Assertions.*;
 
+import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 public class LargeQueryExecutionTests {
+  private Connection connection;
+
+  @BeforeEach
+  void setUp() throws SQLException {
+    connection = getValidJDBCConnection();
+  }
+
+  @AfterEach
+  void cleanUp() throws SQLException {
+    if (connection != null) {
+      connection.close();
+    }
+  }
+
   @Test
   void testQueryYieldingLargeNarrowResultSet() throws SQLException {
     String largeQuerySQL = "SELECT * FROM range(37500000)"; // ~300MB of data
-    ResultSet rs = executeQuery(largeQuerySQL);
+    ResultSet rs = executeQuery(connection, largeQuerySQL);
     int rows = 0;
     while (rs != null && rs.next()) {
       rows++;
@@ -36,7 +53,7 @@ void testQueryYieldingLargeWideResultSet() throws SQLException {
 
     // Create the SQL query
     String query = String.format("SELECT id, %s FROM RANGE(%d)", uuids, rows);
-    ResultSet rs = executeQuery(query);
+    ResultSet rs = executeQuery(connection, query);
     int rowCount = 0;
     while (rs != null && rs.next()) {
       assertEquals(rs.getInt("id"), rowCount, "Expected id to be equal to row number");

src/test/java/com/databricks/jdbc/integration/e2e/MetadataTests.java

@@ -61,17 +61,17 @@ void testResultSetMetadataRetrieval() throws SQLException {
             + "name VARCHAR(255), "
             + "age INT"
             + ");";
-    setupDatabaseTable(tableName, createTableSQL);
+    setupDatabaseTable(connection, tableName, createTableSQL);
 
     String insertSQL =
         "INSERT INTO "
             + getFullyQualifiedTableName(tableName)
             + " (id, name, age) VALUES (1, 'Madhav', 24)";
-    executeSQL(insertSQL);
+    executeSQL(connection, insertSQL);
 
     String query = "SELECT id, name, age FROM " + getFullyQualifiedTableName(tableName);
 
-    ResultSet resultSet = executeQuery(query);
+    ResultSet resultSet = executeQuery(connection, query);
     ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
 
     // Check the number of columns
@@ -106,7 +106,7 @@ void testResultSetMetadataRetrieval() throws SQLException {
           "Column " + i + " should be nullable");
     }
     String SQL = "DROP TABLE IF EXISTS " + getFullyQualifiedTableName(tableName);
-    executeSQL(SQL);
+    executeSQL(connection, SQL);
   }
 
   @Test
@@ -141,7 +141,7 @@ void testTableInformation() throws SQLException {
     String catalog = "main";
     String schemaPattern = "jdbc_test_schema";
     String tableName = "catalog_and_schema_test_table";
-    setupDatabaseTable(tableName);
+    setupDatabaseTable(connection, tableName);
     try (ResultSet tables = metaData.getTables(catalog, schemaPattern, "%", null)) {
       assertTrue(
           tables.next(), "There should be at least one table in the specified catalog and schema");
@@ -150,7 +150,7 @@ void testTableInformation() throws SQLException {
         assertNotNull(fetchedTableName, "Table name should not be null");
       } while (tables.next());
     }
-    deleteTable(tableName);
+    deleteTable(connection, tableName);
   }
 
   @Test
@@ -159,7 +159,7 @@ void testTableInformationExactMatch() throws SQLException {
     String catalog = "main";
     String schemaPattern = "jdbc_test_schema";
     String tableName = "catalog_and_schema_test_table";
-    setupDatabaseTable(tableName);
+    setupDatabaseTable(connection, tableName);
     try (ResultSet tables = metaData.getTables(catalog, schemaPattern, tableName, null)) {
       assertTrue(
           tables.next(), "There should be at least one table in the specified catalog and schema");
@@ -169,6 +169,6 @@ void testTableInformationExactMatch() throws SQLException {
             tableName, fetchedTableName, "Table name should match the specified table name");
       } while (tables.next());
     }
-    deleteTable(tableName);
+    deleteTable(connection, tableName);
   }
 }

src/test/java/com/databricks/jdbc/integration/e2e/UCMetadataTests.java

@@ -30,39 +30,43 @@ static void setUp() throws SQLException {
     // Change connection to actual test warehouse once it supports latest runtime
     connection = getDogfoodJDBCConnection(List.of(List.of("useLegacyMetadata", "0")));
 
-    executeSQL("CREATE CATALOG IF NOT EXISTS " + catA);
-    executeSQL("USE CATALOG " + catA);
-    executeSQL("CREATE DATABASE IF NOT EXISTS " + db1);
-    executeSQL("CREATE DATABASE IF NOT EXISTS " + db2);
+    executeSQL(connection, "CREATE CATALOG IF NOT EXISTS " + catA);
+    executeSQL(connection, "USE CATALOG " + catA);
+    executeSQL(connection, "CREATE DATABASE IF NOT EXISTS " + db1);
+    executeSQL(connection, "CREATE DATABASE IF NOT EXISTS " + db2);
 
-    executeSQL("USE " + db1);
+    executeSQL(connection, "USE " + db1);
     executeSQL(
+        connection,
         "CREATE TABLE IF NOT EXISTS a_1 AS (SELECT 1 AS col_1, '"
             + catA
             + "."
             + db1
             + ".a_1' AS col_2)");
     executeSQL(
+        connection,
         "CREATE TABLE IF NOT EXISTS a_2 AS (SELECT 1 AS col_1, '"
             + catA
             + "."
             + db1
             + ".a_2' AS col_2)");
 
-    executeSQL("USE " + db2);
+    executeSQL(connection, "USE " + db2);
     executeSQL(
+        connection,
         "CREATE TABLE IF NOT EXISTS a_1 AS (SELECT 1 AS col_1, '"
             + catA
             + "."
             + db2
             + ".a_1' AS col_2)");
 
-    executeSQL("USE CATALOG " + mainCatalog);
-    executeSQL("CREATE DATABASE IF NOT EXISTS " + db1);
-    executeSQL("CREATE DATABASE IF NOT EXISTS " + db2);
+    executeSQL(connection, "USE CATALOG " + mainCatalog);
+    executeSQL(connection, "CREATE DATABASE IF NOT EXISTS " + db1);
+    executeSQL(connection, "CREATE DATABASE IF NOT EXISTS " + db2);
 
-    executeSQL("USE " + db1);
+    executeSQL(connection, "USE " + db1);
     executeSQL(
+        connection,
         "CREATE TABLE IF NOT EXISTS "
             + mainDb1Table1
             + " AS (SELECT 1 AS col_1, 'main."
@@ -71,10 +75,11 @@ static void setUp() throws SQLException {
             + mainDb1Table1
             + "' AS col_2)");
 
-    executeSQL("USE CATALOG " + hiveCatalog);
-    executeSQL("CREATE DATABASE IF NOT EXISTS " + db2);
-    executeSQL("USE " + db2);
+    executeSQL(connection, "USE CATALOG " + hiveCatalog);
+    executeSQL(connection, "CREATE DATABASE IF NOT EXISTS " + db2);
+    executeSQL(connection, "USE " + db2);
     executeSQL(
+        connection,
         "CREATE TABLE IF NOT EXISTS "
             + table1
             + " AS (SELECT 1 AS col_1, '"
@@ -85,6 +90,7 @@ static void setUp() throws SQLException {
             + table1
             + "' AS col_2)");
     executeSQL(
+        connection,
         "CREATE TABLE IF NOT EXISTS "
             + table2
             + " AS (SELECT 1 AS col_1, '"
@@ -110,20 +116,20 @@ static void setUp() throws SQLException {
   @AfterAll
   static void cleanUp() throws SQLException {
     // Cleanup
-    executeSQL("USE CATALOG " + catA);
-    executeSQL("DROP DATABASE " + db1 + " CASCADE");
-    executeSQL("DROP DATABASE " + db2 + " CASCADE");
-    executeSQL("DROP DATABASE IF EXISTS default CASCADE");
-    executeSQL("DROP CATALOG " + catA);
+    executeSQL(connection, "USE CATALOG " + catA);
+    executeSQL(connection, "DROP DATABASE " + db1 + " CASCADE");
+    executeSQL(connection, "DROP DATABASE " + db2 + " CASCADE");
+    executeSQL(connection, "DROP DATABASE IF EXISTS default CASCADE");
+    executeSQL(connection, "DROP CATALOG " + catA);
 
     // Cleanup main catalog
-    executeSQL("USE CATALOG " + mainCatalog);
-    executeSQL("DROP DATABASE " + db1 + " CASCADE");
-    executeSQL("DROP DATABASE " + db2 + " CASCADE");
+    executeSQL(connection, "USE CATALOG " + mainCatalog);
+    executeSQL(connection, "DROP DATABASE " + db1 + " CASCADE");
+    executeSQL(connection, "DROP DATABASE " + db2 + " CASCADE");
 
     // Cleanup legacy catalog
-    executeSQL("USE CATALOG " + hiveCatalog);
-    executeSQL("DROP DATABASE " + db2 + " CASCADE");
+    executeSQL(connection, "USE CATALOG " + hiveCatalog);
+    executeSQL(connection, "DROP DATABASE " + db2 + " CASCADE");
 
     if (connection != null) {
       connection.close();
@@ -138,7 +144,7 @@ void testGetCatalogs() throws SQLException {
 
   @Test
   void testGetSchemas() throws SQLException {
-    executeSQL("USE CATALOG hive_metastore");
+    executeSQL(connection, "USE CATALOG hive_metastore");
     ResultSet r = connection.getMetaData().getSchemas("hive_metastore", "%");
     verifyContainsSchemas(
         r,

src/test/java/com/databricks/jdbc/integration/fakeservice/FakeServiceExtension.java

@@ -5,7 +5,6 @@
 import static com.github.tomakehurst.wiremock.common.AbstractFileSource.byFileExtension;
 
 import com.databricks.jdbc.common.DatabricksJdbcConstants.FakeServiceType;
-import com.databricks.jdbc.integration.IntegrationTestUtil;
 import com.github.tomakehurst.wiremock.client.ResponseDefinitionBuilder;
 import com.github.tomakehurst.wiremock.common.BinaryFile;
 import com.github.tomakehurst.wiremock.common.ContentTypes;
@@ -162,7 +161,6 @@ protected void onBeforeAll(WireMockRuntimeInfo wireMockRuntimeInfo, ExtensionCon
     wireMockServerHttpPort = wireMockRuntimeInfo.getHttpPort();
 
     setFakeServiceProperties(wireMockServerHttpPort);
-    IntegrationTestUtil.resetJDBCConnection();
   }
 
   /** {@inheritDoc} */
@@ -182,7 +180,6 @@ protected void onBeforeEach(WireMockRuntimeInfo wireMockRuntimeInfo, ExtensionCo
   @Override
   protected void onAfterEach(WireMockRuntimeInfo wireMockRuntimeInfo, ExtensionContext context)
       throws Exception {
-    IntegrationTestUtil.resetJDBCConnection();
 
     if (fakeServiceMode == FakeServiceMode.RECORD) {
       saveStubMappings(wireMockRuntimeInfo, context);