Skip to content

Commit bee2e4b

Browse files
SNOW-1853752: Clear batch after successful execution (#2024)
1 parent e8e3bac commit bee2e4b

File tree

8 files changed

+346
-15
lines changed

8 files changed

+346
-15
lines changed

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

+13-12
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,8 @@ public abstract class SFBaseSession {
148148

149149
private boolean supportImplicitAsyncQueryTimeout = false;
150150

151+
private boolean clearBatchOnlyAfterSuccessfulExecution = false;
152+
151153
protected SFBaseSession(SFConnectionHandler sfConnectionHandler) {
152154
this.sfConnectionHandler = sfConnectionHandler;
153155
}
@@ -1335,21 +1337,20 @@ public boolean getEnableReturnTimestampWithTimeZone() {
13351337
return enableReturnTimestampWithTimeZone;
13361338
}
13371339

1338-
/**
1339-
* @return True if query timeout should be set on the server side for async queries. False by
1340-
* default.
1341-
*/
1342-
@SnowflakeJdbcInternalApi
1343-
public boolean getSupportImplicitAsyncQueryTimeout() {
1340+
boolean getSupportImplicitAsyncQueryTimeout() {
13441341
return supportImplicitAsyncQueryTimeout;
13451342
}
13461343

1347-
/**
1348-
* @param supportImplicitAsyncQueryTimeout Setting supportImplicitAsyncQueryTimeout to true allows
1349-
* for query timeout to be set on the server side.
1350-
*/
1351-
@SnowflakeJdbcInternalApi
1352-
public void setSupportImplicitAsyncQueryTimeout(boolean supportImplicitAsyncQueryTimeout) {
1344+
void setSupportImplicitAsyncQueryTimeout(boolean supportImplicitAsyncQueryTimeout) {
13531345
this.supportImplicitAsyncQueryTimeout = supportImplicitAsyncQueryTimeout;
13541346
}
1347+
1348+
void setClearBatchOnlyAfterSuccessfulExecution(boolean value) {
1349+
this.clearBatchOnlyAfterSuccessfulExecution = value;
1350+
}
1351+
1352+
@SnowflakeJdbcInternalApi
1353+
public boolean getClearBatchOnlyAfterSuccessfulExecution() {
1354+
return this.clearBatchOnlyAfterSuccessfulExecution;
1355+
}
13551356
}

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

+6
Original file line numberDiff line numberDiff line change
@@ -557,6 +557,12 @@ public void addSFSessionProperty(String propertyName, Object propertyValue) thro
557557
}
558558
break;
559559

560+
case CLEAR_BATCH_ONLY_AFTER_SUCCESSFUL_EXECUTION:
561+
if (propertyValue != null) {
562+
setClearBatchOnlyAfterSuccessfulExecution(getBooleanValue(propertyValue));
563+
}
564+
break;
565+
560566
default:
561567
break;
562568
}

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

+4-1
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,10 @@ public enum SFSessionProperty {
126126
"JAVA_LOGGING_CONSOLE_STD_OUT_THRESHOLD", false, String.class),
127127

128128
SUPPORT_IMPLICIT_ASYNC_QUERY_TIMEOUT(
129-
"SUPPORT_IMPLICIT_ASYNC_QUERY_TIMEOUT", false, Boolean.class);
129+
"SUPPORT_IMPLICIT_ASYNC_QUERY_TIMEOUT", false, Boolean.class),
130+
131+
CLEAR_BATCH_ONLY_AFTER_SUCCESSFUL_EXECUTION(
132+
"CLEAR_BATCH_ONLY_AFTER_SUCCESSFUL_EXECUTION", false, Boolean.class);
130133

131134
// property key in string
132135
private String propertyKey;

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

+10-1
Original file line numberDiff line numberDiff line change
@@ -1013,8 +1013,17 @@ VariableTypeArray executeBatchInternalWithArrayBind(boolean isLong) throws SQLEx
10131013
updateCounts.intArr = executeBatchInternal(false).intArr;
10141014
}
10151015
}
1016+
if (this.getSFBaseStatement()
1017+
.getSFBaseSession()
1018+
.getClearBatchOnlyAfterSuccessfulExecution()) {
1019+
clearBatch();
1020+
}
10161021
} finally {
1017-
this.clearBatch();
1022+
if (!this.getSFBaseStatement()
1023+
.getSFBaseSession()
1024+
.getClearBatchOnlyAfterSuccessfulExecution()) {
1025+
clearBatch();
1026+
}
10181027
}
10191028

10201029
return updateCounts;

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

+3-1
Original file line numberDiff line numberDiff line change
@@ -556,7 +556,9 @@ VariableTypeArray executeBatchInternal(boolean isLong) throws SQLException {
556556
updateCounts.intArr,
557557
exceptionReturned);
558558
}
559-
559+
if (this.getSFBaseStatement().getSFBaseSession().getClearBatchOnlyAfterSuccessfulExecution()) {
560+
clearBatch();
561+
}
560562
return updateCounts;
561563
}
562564

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

+16
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,10 @@
1414
import java.sql.Statement;
1515
import java.util.Arrays;
1616
import java.util.List;
17+
import java.util.Random;
18+
import java.util.UUID;
1719
import java.util.regex.Pattern;
20+
import java.util.stream.Collectors;
1821
import net.snowflake.client.core.SFException;
1922
import net.snowflake.client.jdbc.SnowflakeUtil;
2023
import net.snowflake.client.log.SFLogger;
@@ -154,4 +157,17 @@ public static void expectSnowflakeLoggedFeatureNotSupportedException(MethodRaise
154157
public static void assertEqualsIgnoringWhitespace(String expected, String actual) {
155158
assertEquals(expected.replaceAll("\\s+", ""), actual.replaceAll("\\s+", ""));
156159
}
160+
161+
public static String randomTableName(String jiraId) {
162+
return ("TEST_" + (jiraId != null ? jiraId : "") + "_" + UUID.randomUUID())
163+
.replaceAll("-", "_");
164+
}
165+
166+
public static List<Integer> randomIntList(int length, int modulo) {
167+
return new Random()
168+
.ints()
169+
.limit(length)
170+
.mapToObj(i -> Math.abs(i) % modulo)
171+
.collect(Collectors.toList());
172+
}
157173
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
package net.snowflake.client.jdbc;
2+
3+
import static net.snowflake.client.TestUtil.randomIntList;
4+
import static org.junit.jupiter.api.Assertions.assertEquals;
5+
import static org.junit.jupiter.api.Assertions.assertThrows;
6+
7+
import java.sql.BatchUpdateException;
8+
import java.sql.Connection;
9+
import java.sql.SQLException;
10+
import java.sql.Statement;
11+
import java.util.Arrays;
12+
import java.util.List;
13+
import net.snowflake.client.TestUtil;
14+
import net.snowflake.client.category.TestTags;
15+
import org.junit.jupiter.api.Tag;
16+
import org.junit.jupiter.api.Test;
17+
18+
@Tag(TestTags.STATEMENT)
19+
public class BatchExecutionIT extends BaseJDBCTest {
20+
@Test
21+
public void testClearingBatchAfterStatementExecution() throws SQLException {
22+
String tableName = TestUtil.randomTableName("SNOW-1853752");
23+
int itemsInBatch = 3;
24+
try (Connection connection = getConnection();
25+
Statement statement = connection.createStatement()) {
26+
statement.execute(
27+
String.format("CREATE OR REPLACE TABLE %s(id int, s varchar(2))", tableName));
28+
List<ThrowingCallable<Integer, SQLException>> executeMethods =
29+
Arrays.asList(
30+
() -> statement.executeBatch().length, () -> statement.executeLargeBatch().length);
31+
for (ThrowingCallable<Integer, SQLException> executeMethod : executeMethods) {
32+
for (int i : randomIntList(itemsInBatch, 10)) {
33+
statement.addBatch(
34+
String.format("INSERT INTO %s(id, s) VALUES (%d, 's%d')", tableName, i, i));
35+
}
36+
assertEquals(itemsInBatch, executeMethod.call());
37+
// default behaviour - batch is not cleared
38+
assertEquals(itemsInBatch, executeMethod.call());
39+
statement.clearBatch();
40+
for (int i : randomIntList(itemsInBatch, 10)) {
41+
statement.addBatch(
42+
String.format(
43+
"INSERT INTO %s(id, s) VALUES (%d, 'longer string %d')", tableName, i, i));
44+
}
45+
assertThrows(BatchUpdateException.class, executeMethod::call);
46+
// second call should also fail since batch should not be cleared
47+
assertThrows(BatchUpdateException.class, executeMethod::call);
48+
// clear batch for next execution in loop
49+
statement.clearBatch();
50+
}
51+
}
52+
}
53+
54+
/**
55+
* ThrowingCallable is defined here since is not available in OldDriverTests from main package.
56+
* Can be removed when OldDriver version is set to >=3.15.1
57+
*/
58+
@FunctionalInterface
59+
interface ThrowingCallable<A, T extends Throwable> {
60+
A call() throws T;
61+
}
62+
}

0 commit comments

Comments
 (0)