SNOW-780605: Better handle resource closing in testHTAPOptimizations (#…

…1500)
snowflakedb · Aug 28, 2023 · 1e92087 · 1e92087
1 parent 8759eb1
commit 1e92087
Showing 1 changed file with 60 additions and 48 deletions.
diff --git a/src/test/java/net/snowflake/client/jdbc/SnowflakeDriverLatestIT.java b/src/test/java/net/snowflake/client/jdbc/SnowflakeDriverLatestIT.java
@@ -1570,54 +1570,66 @@ public void testUploadWithGCSDownscopedCredentialWithoutConnection() throws Thro
   @Test
   @ConditionalIgnoreRule.ConditionalIgnore(condition = RunningOnGithubAction.class)
   public void testHTAPOptimizations() throws SQLException {
-    // Set the HTAP test parameter to true
-    try (Connection con = getSnowflakeAdminConnection()) {
-      Statement statement = con.createStatement();
-      statement.execute(
-          "alter account "
-              + TestUtil.systemGetEnv("SNOWFLAKE_TEST_ACCOUNT")
-              + " set ENABLE_SNOW_654741_FOR_TESTING=true");
-    }
-    // Create a normal connection and assert that database, schema, and warehouse have expected
-    // values
-    Connection con = getConnection();
-    SFSession session = con.unwrap(SnowflakeConnectionV1.class).getSfSession();
-    assertTrue(TestUtil.systemGetEnv("SNOWFLAKE_TEST_SCHEMA").equalsIgnoreCase(con.getSchema()));
-    assertTrue(TestUtil.systemGetEnv("SNOWFLAKE_TEST_DATABASE").equalsIgnoreCase(con.getCatalog()));
-    assertTrue(
-        TestUtil.systemGetEnv("SNOWFLAKE_TEST_WAREHOUSE").equalsIgnoreCase(session.getWarehouse()));
-    Statement statement = con.createStatement();
-    // Set TIMESTAMP_OUTPUT_FORMAT (which is a session parameter) to check its value later
-    statement.execute("alter session set TIMESTAMP_OUTPUT_FORMAT='YYYY-MM-DD HH24:MI:SS.FFTZH'");
-    statement.execute("create or replace table testtable1 (cola string, colb int)");
-    statement.execute("insert into testtable1 values ('row1', 1), ('row2', 2), ('row3', 3)");
-    ResultSet rs = statement.executeQuery("select * from testtable1");
-    assertEquals(3, getSizeOfResultSet(rs));
-    // Assert database, schema, and warehouse have the same values as before even though the select
-    // statement will
-    // return no parameters or metadata
-    assertTrue(TestUtil.systemGetEnv("SNOWFLAKE_TEST_SCHEMA").equalsIgnoreCase(con.getSchema()));
-    assertTrue(TestUtil.systemGetEnv("SNOWFLAKE_TEST_DATABASE").equalsIgnoreCase(con.getCatalog()));
-    assertTrue(
-        TestUtil.systemGetEnv("SNOWFLAKE_TEST_WAREHOUSE").equalsIgnoreCase(session.getWarehouse()));
-    // Assert session parameter TIMESTAMP_OUTPUT_FORMAT has the same value as before the select
-    // statement
-    assertEquals(
-        "YYYY-MM-DD HH24:MI:SS.FFTZH",
-        session.getCommonParameters().get("TIMESTAMP_OUTPUT_FORMAT"));
-    // cleanup
-    statement.execute("alter session unset TIMESTAMP_OUTPUT_FORMAT");
-    statement.execute("drop table testtable1");
-    rs.close();
-    statement.close();
-    con.close();
-    // cleanup
-    try (Connection con2 = getSnowflakeAdminConnection()) {
-      statement = con2.createStatement();
-      statement.execute(
-          "alter account "
-              + TestUtil.systemGetEnv("SNOWFLAKE_TEST_ACCOUNT")
-              + " unset ENABLE_SNOW_654741_FOR_TESTING");
+    try {
+      // Set the HTAP test parameter to true
+      try (Connection con = getSnowflakeAdminConnection();
+          Statement statement = con.createStatement()) {
+        statement.execute(
+            "alter account "
+                + TestUtil.systemGetEnv("SNOWFLAKE_TEST_ACCOUNT")
+                + " set ENABLE_SNOW_654741_FOR_TESTING=true");
+      }
+      // Create a normal connection and assert that database, schema, and warehouse have expected
+      // values
+      try (Connection con = getConnection()) {
+        SFSession session = con.unwrap(SnowflakeConnectionV1.class).getSfSession();
+        assertTrue(
+            TestUtil.systemGetEnv("SNOWFLAKE_TEST_SCHEMA").equalsIgnoreCase(con.getSchema()));
+        assertTrue(
+            TestUtil.systemGetEnv("SNOWFLAKE_TEST_DATABASE").equalsIgnoreCase(con.getCatalog()));
+        assertTrue(
+            TestUtil.systemGetEnv("SNOWFLAKE_TEST_WAREHOUSE")
+                .equalsIgnoreCase(session.getWarehouse()));
+        try (Statement statement = con.createStatement()) {
+          // Set TIMESTAMP_OUTPUT_FORMAT (which is a session parameter) to check its value later
+          try {
+            statement.execute(
+                "alter session set TIMESTAMP_OUTPUT_FORMAT='YYYY-MM-DD HH24:MI:SS.FFTZH'");
+            statement.execute("create or replace table testtable1 (cola string, colb int)");
+            statement.execute(
+                "insert into testtable1 values ('row1', 1), ('row2', 2), ('row3', 3)");
+            try (ResultSet rs = statement.executeQuery("select * from testtable1")) {
+              assertEquals(3, getSizeOfResultSet(rs));
+              // Assert database, schema, and warehouse have the same values as before even though
+              // the select statement will return no parameters or metadata
+              assertTrue(
+                  TestUtil.systemGetEnv("SNOWFLAKE_TEST_SCHEMA").equalsIgnoreCase(con.getSchema()));
+              assertTrue(
+                  TestUtil.systemGetEnv("SNOWFLAKE_TEST_DATABASE")
+                      .equalsIgnoreCase(con.getCatalog()));
+              assertTrue(
+                  TestUtil.systemGetEnv("SNOWFLAKE_TEST_WAREHOUSE")
+                      .equalsIgnoreCase(session.getWarehouse()));
+              // Assert session parameter TIMESTAMP_OUTPUT_FORMAT has the same value as before the
+              // select statement
+              assertEquals(
+                  "YYYY-MM-DD HH24:MI:SS.FFTZH",
+                  session.getCommonParameters().get("TIMESTAMP_OUTPUT_FORMAT"));
+            }
+          } finally {
+            statement.execute("alter session unset TIMESTAMP_OUTPUT_FORMAT");
+            statement.execute("drop table if exists testtable1");
+          }
+        }
+      }
+    } finally {
+      try (Connection con = getSnowflakeAdminConnection();
+          Statement statement = con.createStatement()) {
+        statement.execute(
+            "alter account "
+                + TestUtil.systemGetEnv("SNOWFLAKE_TEST_ACCOUNT")
+                + " unset ENABLE_SNOW_654741_FOR_TESTING");
+      }
     }
   }