SNOW-1853752: Clear batch after successful execution (#2024)

Author: sfc-gh-dprzybysz

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

@@ -148,6 +148,8 @@ public abstract class SFBaseSession {
 
   private boolean supportImplicitAsyncQueryTimeout = false;
 
+  private boolean clearBatchOnlyAfterSuccessfulExecution = false;
+
   protected SFBaseSession(SFConnectionHandler sfConnectionHandler) {
     this.sfConnectionHandler = sfConnectionHandler;
   }
@@ -1335,21 +1337,20 @@ public boolean getEnableReturnTimestampWithTimeZone() {
     return enableReturnTimestampWithTimeZone;
   }
 
-  /**
-   * @return True if query timeout should be set on the server side for async queries. False by
-   *     default.
-   */
-  @SnowflakeJdbcInternalApi
-  public boolean getSupportImplicitAsyncQueryTimeout() {
+  boolean getSupportImplicitAsyncQueryTimeout() {
     return supportImplicitAsyncQueryTimeout;
   }
 
-  /**
-   * @param supportImplicitAsyncQueryTimeout Setting supportImplicitAsyncQueryTimeout to true allows
-   *     for query timeout to be set on the server side.
-   */
-  @SnowflakeJdbcInternalApi
-  public void setSupportImplicitAsyncQueryTimeout(boolean supportImplicitAsyncQueryTimeout) {
+  void setSupportImplicitAsyncQueryTimeout(boolean supportImplicitAsyncQueryTimeout) {
     this.supportImplicitAsyncQueryTimeout = supportImplicitAsyncQueryTimeout;
   }
+
+  void setClearBatchOnlyAfterSuccessfulExecution(boolean value) {
+    this.clearBatchOnlyAfterSuccessfulExecution = value;
+  }
+
+  @SnowflakeJdbcInternalApi
+  public boolean getClearBatchOnlyAfterSuccessfulExecution() {
+    return this.clearBatchOnlyAfterSuccessfulExecution;
+  }
 }

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

@@ -557,6 +557,12 @@ public void addSFSessionProperty(String propertyName, Object propertyValue) thro
           }
           break;
 
+        case CLEAR_BATCH_ONLY_AFTER_SUCCESSFUL_EXECUTION:
+          if (propertyValue != null) {
+            setClearBatchOnlyAfterSuccessfulExecution(getBooleanValue(propertyValue));
+          }
+          break;
+
         default:
           break;
       }

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

@@ -126,7 +126,10 @@ public enum SFSessionProperty {
       "JAVA_LOGGING_CONSOLE_STD_OUT_THRESHOLD", false, String.class),
 
   SUPPORT_IMPLICIT_ASYNC_QUERY_TIMEOUT(
-      "SUPPORT_IMPLICIT_ASYNC_QUERY_TIMEOUT", false, Boolean.class);
+      "SUPPORT_IMPLICIT_ASYNC_QUERY_TIMEOUT", false, Boolean.class),
+
+  CLEAR_BATCH_ONLY_AFTER_SUCCESSFUL_EXECUTION(
+      "CLEAR_BATCH_ONLY_AFTER_SUCCESSFUL_EXECUTION", false, Boolean.class);
 
   // property key in string
   private String propertyKey;

src/main/java/net/snowflake/client/jdbc/SnowflakePreparedStatementV1.java

@@ -1013,8 +1013,17 @@ VariableTypeArray executeBatchInternalWithArrayBind(boolean isLong) throws SQLEx
           updateCounts.intArr = executeBatchInternal(false).intArr;
         }
       }
+      if (this.getSFBaseStatement()
+          .getSFBaseSession()
+          .getClearBatchOnlyAfterSuccessfulExecution()) {
+        clearBatch();
+      }
     } finally {
-      this.clearBatch();
+      if (!this.getSFBaseStatement()
+          .getSFBaseSession()
+          .getClearBatchOnlyAfterSuccessfulExecution()) {
+        clearBatch();
+      }
     }
 
     return updateCounts;

src/main/java/net/snowflake/client/jdbc/SnowflakeStatementV1.java

@@ -556,7 +556,9 @@ VariableTypeArray executeBatchInternal(boolean isLong) throws SQLException {
           updateCounts.intArr,
           exceptionReturned);
     }
-
+    if (this.getSFBaseStatement().getSFBaseSession().getClearBatchOnlyAfterSuccessfulExecution()) {
+      clearBatch();
+    }
     return updateCounts;
   }
 

src/test/java/net/snowflake/client/TestUtil.java

@@ -14,7 +14,10 @@
 import java.sql.Statement;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Random;
+import java.util.UUID;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 import net.snowflake.client.core.SFException;
 import net.snowflake.client.jdbc.SnowflakeUtil;
 import net.snowflake.client.log.SFLogger;
@@ -154,4 +157,17 @@ public static void expectSnowflakeLoggedFeatureNotSupportedException(MethodRaise
   public static void assertEqualsIgnoringWhitespace(String expected, String actual) {
     assertEquals(expected.replaceAll("\\s+", ""), actual.replaceAll("\\s+", ""));
   }
+
+  public static String randomTableName(String jiraId) {
+    return ("TEST_" + (jiraId != null ? jiraId : "") + "_" + UUID.randomUUID())
+        .replaceAll("-", "_");
+  }
+
+  public static List<Integer> randomIntList(int length, int modulo) {
+    return new Random()
+        .ints()
+        .limit(length)
+        .mapToObj(i -> Math.abs(i) % modulo)
+        .collect(Collectors.toList());
+  }
 }

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

@@ -0,0 +1,62 @@
+package net.snowflake.client.jdbc;
+
+import static net.snowflake.client.TestUtil.randomIntList;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import java.sql.BatchUpdateException;
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.Arrays;
+import java.util.List;
+import net.snowflake.client.TestUtil;
+import net.snowflake.client.category.TestTags;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+
+@Tag(TestTags.STATEMENT)
+public class BatchExecutionIT extends BaseJDBCTest {
+  @Test
+  public void testClearingBatchAfterStatementExecution() throws SQLException {
+    String tableName = TestUtil.randomTableName("SNOW-1853752");
+    int itemsInBatch = 3;
+    try (Connection connection = getConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute(
+          String.format("CREATE OR REPLACE TABLE %s(id int, s varchar(2))", tableName));
+      List<ThrowingCallable<Integer, SQLException>> executeMethods =
+          Arrays.asList(
+              () -> statement.executeBatch().length, () -> statement.executeLargeBatch().length);
+      for (ThrowingCallable<Integer, SQLException> executeMethod : executeMethods) {
+        for (int i : randomIntList(itemsInBatch, 10)) {
+          statement.addBatch(
+              String.format("INSERT INTO %s(id, s) VALUES (%d, 's%d')", tableName, i, i));
+        }
+        assertEquals(itemsInBatch, executeMethod.call());
+        // default behaviour - batch is not cleared
+        assertEquals(itemsInBatch, executeMethod.call());
+        statement.clearBatch();
+        for (int i : randomIntList(itemsInBatch, 10)) {
+          statement.addBatch(
+              String.format(
+                  "INSERT INTO %s(id, s) VALUES (%d, 'longer string %d')", tableName, i, i));
+        }
+        assertThrows(BatchUpdateException.class, executeMethod::call);
+        // second call should also fail since batch should not be cleared
+        assertThrows(BatchUpdateException.class, executeMethod::call);
+        // clear batch for next execution in loop
+        statement.clearBatch();
+      }
+    }
+  }
+
+  /**
+   * ThrowingCallable is defined here since is not available in OldDriverTests from main package.
+   * Can be removed when OldDriver version is set to >=3.15.1
+   */
+  @FunctionalInterface
+  interface ThrowingCallable<A, T extends Throwable> {
+    A call() throws T;
+  }
+}