diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java b/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java index b2de2c63ad..7511efff06 100644 --- a/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java +++ b/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java @@ -7537,7 +7537,7 @@ protected void setInterruptsEnabled(boolean interruptsEnabled) { // Flag set to indicate that an interrupt has happened. private volatile boolean wasInterrupted = false; - private boolean wasInterrupted() { + boolean wasInterrupted() { return wasInterrupted; } diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java index 17bd4a5fe9..368876b538 100644 --- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java +++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java @@ -2804,6 +2804,21 @@ final void doExecutePreparedStatementBatch(PrepStmtBatchExecCmd batchCommand) th for (int attempt = 1; attempt <= 2; ++attempt) { try { + // If the command was interrupted, that means the TDS.PKT_CANCEL_REQ was sent to the server. + // Since the cancellation request was sent, stop processing the batch query and process the + // cancellation request and then return. + // + // Otherwise, if we do continue processing the batch query, in the case where a query requires + // prepexec/sp_prepare, the TDS request for prepexec/sp_prepare will be sent regardless of + // query cancellation. This will cause a TDS token error in the post processing when we + // close the query. + if (batchCommand.wasInterrupted()) { + ensureExecuteResultsReader(batchCommand.startResponse(getIsResponseBufferingAdaptive())); + startResults(); + getNextResult(true); + return; + } + // Re-use handle if available, requires parameter definitions which are not available until here. if (reuseCachedHandle(hasNewTypeDefinitions, 1 < attempt)) { hasNewTypeDefinitions = false; diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java index 1bbb293844..cc652f4c0e 100644 --- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java +++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java @@ -48,7 +48,7 @@ protected Object[][] getContents() { {"R_noServerResponse", "SQL Server did not return a response. The connection has been closed."}, {"R_truncatedServerResponse", "SQL Server returned an incomplete response. The connection has been closed."}, {"R_queryTimedOut", "The query has timed out."}, - {"R_queryCancelled", "The query was canceled."}, + {"R_queryCancelled", "The query was cancelled."}, {"R_errorReadingStream", "An error occurred while reading the value from the stream object. Error: \"{0}\""}, {"R_streamReadReturnedInvalidValue", "The stream read operation returned an invalid value for the amount of data read."}, {"R_mismatchedStreamLength", "The stream value is not the specified length. The specified length was {0}, the actual length is {1}."}, diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java b/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java index da6773116c..4f7d5fa69d 100644 --- a/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java +++ b/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java @@ -161,7 +161,6 @@ protected Object[][] getContents() { {"R_cancellationFailed", "Cancellation failed."}, {"R_executionNotTimeout", "Execution did not timeout."}, {"R_executionTooLong", "Execution took too long."}, {"R_executionNotLong", "Execution did not take long enough."}, - {"R_queryCancelled", "The query was canceled."}, {"R_statementShouldBeClosed", "statement should be closed since resultset is closed."}, {"R_statementShouldBeOpened", "statement should be opened since resultset is opened."}, {"R_shouldBeWrapper", "{0} should be a wrapper for {1}."}, @@ -201,5 +200,6 @@ protected Object[][] getContents() { {"R_objectNullOrEmpty", "The {0} is null or empty."}, {"R_cekDecryptionFailed", "Failed to decrypt a column encryption key using key store provider: {0}."}, {"R_connectTimedOut", "connect timed out"}, + {"R_queryCancelled", "The query was cancelled."}, {"R_sessionKilled", "Cannot continue the execution because the session is in the kill state"}}; } diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/BatchExecutionTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/BatchExecutionTest.java index 8cb258a03a..90a3b0a9ac 100644 --- a/src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/BatchExecutionTest.java +++ b/src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/BatchExecutionTest.java @@ -5,6 +5,7 @@ package com.microsoft.sqlserver.jdbc.unit.statement; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.fail; import java.lang.reflect.Field; @@ -60,6 +61,51 @@ public void testBatchExceptionAEOn() throws Exception { testExecuteBatch1UseBulkCopyAPI(); } + @Test + public void testBatchStatementCancellation() throws Exception { + try (Connection connection = PrepUtil.getConnection(connectionString)) { + connection.setAutoCommit(false); + + try (PreparedStatement statement = connection.prepareStatement( + "if object_id('test_table') is not null drop table test_table")) { + statement.execute(); + } + connection.commit(); + + try (PreparedStatement statement = connection.prepareStatement( + "create table test_table (column_name bit)")) { + statement.execute(); + } + connection.commit(); + + for (long delayInMilliseconds : new long[] { 1, 2, 4, 8, 16, 32, 64, 128 }) { + for (int numberOfCommands : new int[] { 1, 2, 4, 8, 16, 32, 64 }) { + int parameterCount = 512; + + try (PreparedStatement statement = connection.prepareStatement( + "insert into test_table values (?)" + repeat(",(?)", parameterCount - 1))) { + + for (int i = 0; i < numberOfCommands; i++) { + for (int j = 0; j < parameterCount; j++) { + statement.setBoolean(j + 1, true); + } + statement.addBatch(); + } + + Thread cancelThread = cancelAsync(statement, delayInMilliseconds); + try { + statement.executeBatch(); + } catch (SQLException e) { + assertEquals(TestResource.getResource("R_queryCancelled"), e.getMessage()); + } + cancelThread.join(); + } + connection.commit(); + } + } + } + } + /** * Get a PreparedStatement object and call the addBatch() method with 3 SQL statements and call the executeBatch() * method and it should return array of Integer values of length 3 @@ -231,6 +277,29 @@ private void modifyConnectionForBulkCopyAPI(SQLServerConnection con) throws Exce con.setUseBulkCopyForBatchInsert(true); } + private static String repeat(String string, int count) { + StringBuilder sb = new StringBuilder(); + for (int i = 0; i < count; i++) { + sb.append(string); + } + return sb.toString(); + } + + private static Thread cancelAsync(Statement statement, long delayInMilliseconds) { + Thread thread = new Thread(() -> { + try { + Thread.sleep(delayInMilliseconds); + statement.cancel(); + } catch (SQLException | InterruptedException e) { + // does not/must not happen + e.printStackTrace(); + throw new IllegalStateException(e); + } + }); + thread.start(); + return thread; + } + @BeforeAll public static void testSetup() throws TestAbortedException, Exception { connectionString = TestUtils.addOrOverrideProperty(connectionString,"trustServerCertificate", "true");