Snow 2234383 test coverage for multistatement jdbc

src/main/java/net/snowflake/client/core/SFResultSetFactory.java

@@ -5,12 +5,15 @@
 import net.snowflake.client.jdbc.ErrorCode;
 import net.snowflake.client.jdbc.SnowflakeResultSetSerializableV1;
 import net.snowflake.client.jdbc.SnowflakeSQLLoggedException;
+import net.snowflake.client.log.SFLogger;
+import net.snowflake.client.log.SFLoggerFactory;
 
 /**
  * Factory class to create SFBaseResultSet class. Depending on result format, different instance
  * will be created
  */
 class SFResultSetFactory {
+  private static final SFLogger logger = SFLoggerFactory.getLogger(SFResultSetFactory.class);
   /**
    * Factory class used to generate ResultSet object according to query result format
    *
@@ -34,6 +37,7 @@ static SFBaseResultSet getResultSet(
     execTimeData.setCreateResultSetStart();
     switch (resultSetSerializable.getQueryResultFormat()) {
       case ARROW:
+        logger.debug("Query result received in ARROW format. Processing with SFArrowResultSet.");
         rs =
             new SFArrowResultSet(
                 resultSetSerializable, statement.getSFBaseSession(), statement, sortResult);

src/test/java/net/snowflake/client/jdbc/MultiStatementLatestIT.java

@@ -5,14 +5,24 @@
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
 
+import java.io.ByteArrayOutputStream;
+import java.io.PrintStream;
+import java.math.BigDecimal;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.logging.ConsoleHandler;
+import java.util.logging.Handler;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 import net.snowflake.client.category.TestTags;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Tag;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 
 /**
  * MultiStatement integration tests for the latest JDBC driver. This doesn't work for the oldest
@@ -314,4 +324,104 @@ public void testAnonymousBlocksUse() throws SQLException {
       }
     }
   }
+
+  @ParameterizedTest
+  @ValueSource(strings = {"arrow", "json"})
+  void testMultiStatementResultFormat(String queryResultFormat) throws Exception {
+    // Setup a memory buffer for JUL logs
+    ByteArrayOutputStream logCapture = new ByteArrayOutputStream();
+    PrintStream logStream = new PrintStream(logCapture);
+
+    Logger sfLogger = Logger.getLogger("net.snowflake");
+    sfLogger.setLevel(Level.ALL);
+    sfLogger.setUseParentHandlers(false);
+
+    Handler consoleHandler =
+        new ConsoleHandler() {
+          @Override
+          public synchronized void publish(java.util.logging.LogRecord record) {
+            logStream.println(record.getLevel() + ": " + record.getMessage());
+          }
+        };
+    consoleHandler.setLevel(Level.ALL);
+    sfLogger.addHandler(consoleHandler);
+
+    try (Statement stmt = connection.createStatement()) {
+      try {
+        stmt.execute("alter session set jdbc_query_result_format = '" + queryResultFormat + "'");
+        stmt.execute("ALTER SESSION SET ENABLE_FIX_1758055_ADD_ARROW_SUPPORT_FOR_MULTI_STMTS=TRUE");
+      } catch (SQLException ex) {
+        /*
+         * Ingore failure since the test user might not be able to change the parameter.
+         * In such case assume the parameter has been enabled on the test account.
+         */
+        if (!ex.getMessage()
+            .contains("invalid parameter 'ENABLE_FIX_1758055_ADD_ARROW_SUPPORT_FOR_MULTI_STMTS'")) {
+          fail();
+        }
+      }
+
+      stmt.unwrap(SnowflakeStatement.class).setParameter("MULTI_STATEMENT_COUNT", 4);
+
+      String multiStmtQuery =
+          "select TO_DOUBLE('0.123456789123456789');\n"
+              + "select 456;\n"
+              + "select 789;\n"
+              + "select '000';";
+
+      stmt.execute(multiStmtQuery);
+
+      int resultSetIndex = 0;
+      boolean hasResults = true;
+
+      while (hasResults) {
+        ResultSet rs = stmt.getResultSet();
+        assertNotNull(rs);
+
+        switch (resultSetIndex) {
+          case 0:
+            assertTrue(rs.next());
+            if ("arrow".equals(queryResultFormat)) {
+              assertEquals(0.12345678912345678d, rs.getDouble(1), 0.0);
+            } else {
+              assertEquals(new BigDecimal("0.1234567891"), rs.getBigDecimal(1));
+            }
+            break;
+          case 1:
+            assertTrue(rs.next());
+            assertEquals(456, rs.getInt(1));
+            break;
+          case 2:
+            assertTrue(rs.next());
+            assertEquals(789, rs.getInt(1));
+            break;
+          case 3:
+            assertTrue(rs.next());
+            assertEquals("000", rs.getString(1));
+            break;
+          default:
+            fail("Unexpected extra result set");
+        }
+
+        rs.close();
+        hasResults = stmt.getMoreResults();
+        resultSetIndex++;
+      }
+    }
+
+    try {
+      consoleHandler.flush();
+      String logs = logCapture.toString();
+
+      if ("arrow".equals(queryResultFormat) && logs != null) {
+        assertTrue(
+            logs.contains(
+                "Query result received in ARROW format. Processing with SFArrowResultSet."),
+            "Expected Arrow logs but got: \n" + logs);
+      }
+    } finally {
+      sfLogger.removeHandler(consoleHandler);
+      consoleHandler.close();
+    }
+  }
 }