fix: executeBatch did not respect statement timeout (#871)

The Statement#executeBatch() method did not respect the statement
timeout that had been set on a statement. This meant that a batch like
the following would always use the default timeout that is set in the
Gapic generated client, and ignore the timeout set in
Statement#setTimeout(int):

statement.setTimeout(120);
statement.addBatch("insert into foo (id) values (1)");
statement.executeBatch(); // This ignored the statement timeout

Fixes #870

Author: olavloite

pom.xml

@@ -220,7 +220,10 @@
         <groupId>org.apache.maven.plugins</groupId>
         <artifactId>maven-surefire-plugin</artifactId>
         <configuration>
-          <exclude>com.google.cloud.spanner.IntegrationTest</exclude>
+          <excludes>
+            <exclude>com.google.cloud.spanner.jdbc.it.**</exclude>
+            <exclude>com.google.cloud.spanner.jdbc.JdbcStatementTimeoutTest</exclude>
+          </excludes>
           <reportNameSuffix>sponge_log</reportNameSuffix>
         </configuration>
       </plugin>

src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java

@@ -104,7 +104,7 @@ private TimeUnit getAppropriateTimeUnit() {
    * specified for a JDBC statement, and then after executing the JDBC statement setting the timeout
    * on the Spanner {@link Connection} again.
    */
-  private static class StatementTimeout {
+  static class StatementTimeout {
     private final long timeout;
     private final TimeUnit unit;
 
@@ -123,7 +123,7 @@ private StatementTimeout(long timeout, TimeUnit unit) {
    * {@link Statement} and returns the original timeout of the Spanner {@link Connection} so it can
    * be reset after the execution of a statement
    */
-  private StatementTimeout setTemporaryStatementTimeout() throws SQLException {
+  StatementTimeout setTemporaryStatementTimeout() throws SQLException {
     StatementTimeout originalTimeout = null;
     if (getQueryTimeout() > 0) {
       if (connection.getSpannerConnection().hasStatementTimeout()) {
@@ -140,7 +140,7 @@ private StatementTimeout setTemporaryStatementTimeout() throws SQLException {
    * Resets the statement timeout of the Spanner {@link Connection} after a JDBC {@link Statement}
    * has been executed.
    */
-  private void resetStatementTimeout(StatementTimeout originalTimeout) throws SQLException {
+  void resetStatementTimeout(StatementTimeout originalTimeout) throws SQLException {
     if (getQueryTimeout() > 0) {
       if (originalTimeout == null) {
         connection.getSpannerConnection().clearStatementTimeout();

src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java

@@ -265,6 +265,7 @@ public long[] executeLargeBatch() throws SQLException {
   private long[] executeBatch(boolean large) throws SQLException {
     checkClosed();
     checkConnectionHasNoActiveBatch();
+    StatementTimeout originalTimeout = setTemporaryStatementTimeout();
     try {
       switch (this.currentBatchType) {
         case DML:
@@ -310,6 +311,7 @@ private long[] executeBatch(boolean large) throws SQLException {
               String.format("Unknown batch type: %s", this.currentBatchType.name()));
       }
     } finally {
+      resetStatementTimeout(originalTimeout);
       batchedStatements.clear();
       this.currentBatchType = BatchType.NONE;
     }

src/test/java/com/google/cloud/spanner/jdbc/JdbcStatementTimeoutTest.java

@@ -0,0 +1,121 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.spanner.jdbc;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
+
+import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
+import com.google.cloud.spanner.connection.AbstractMockServerTest;
+import com.google.cloud.spanner.jdbc.JdbcSqlExceptionFactory.JdbcSqlTimeoutException;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * Tests setting a statement timeout. This test is by default not included in unit test runs, as the
+ * minimum timeout value in JDBC is 1 second, which again makes this test relatively slow.
+ */
+@RunWith(JUnit4.class)
+public class JdbcStatementTimeoutTest extends AbstractMockServerTest {
+
+  @Test
+  public void testExecuteTimeout() throws SQLException {
+    try (java.sql.Connection connection = createJdbcConnection()) {
+      try (Statement statement = connection.createStatement()) {
+        // First verify that execute does not time out by default.
+        assertFalse(statement.execute(INSERT_STATEMENT.getSql()));
+        int result = statement.getUpdateCount();
+        assertEquals(1, result);
+
+        // Simulate that executeSql takes 2 seconds and set a statement timeout of 1 second.
+        mockSpanner.setExecuteSqlExecutionTime(
+            SimulatedExecutionTime.ofMinimumAndRandomTime(2000, 0));
+        statement.setQueryTimeout(1);
+        assertThrows(
+            JdbcSqlTimeoutException.class, () -> statement.execute(INSERT_STATEMENT.getSql()));
+      }
+    }
+  }
+
+  @Test
+  public void testExecuteQueryTimeout() throws SQLException {
+    try (java.sql.Connection connection = createJdbcConnection()) {
+      try (Statement statement = connection.createStatement()) {
+        // First verify that executeQuery does not time out by default.
+        try (ResultSet resultSet = statement.executeQuery(SELECT_RANDOM_STATEMENT.getSql())) {
+          int count = 0;
+          while (resultSet.next()) {
+            count++;
+          }
+          assertEquals(RANDOM_RESULT_SET_ROW_COUNT, count);
+        }
+
+        // Simulate that executeStreamingSql takes 2 seconds and set a statement timeout of 1
+        // second.
+        mockSpanner.setExecuteStreamingSqlExecutionTime(
+            SimulatedExecutionTime.ofMinimumAndRandomTime(2000, 0));
+        statement.setQueryTimeout(1);
+        assertThrows(
+            JdbcSqlTimeoutException.class,
+            () -> statement.executeQuery(SELECT_RANDOM_STATEMENT.getSql()));
+      }
+    }
+  }
+
+  @Test
+  public void testExecuteUpdateTimeout() throws SQLException {
+    try (java.sql.Connection connection = createJdbcConnection()) {
+      try (Statement statement = connection.createStatement()) {
+        // First verify that executeUpdate does not time out by default.
+        assertEquals(1, statement.executeUpdate(INSERT_STATEMENT.getSql()));
+
+        // Simulate that executeSql takes 2 seconds and set a statement timeout of 1 second.
+        mockSpanner.setExecuteSqlExecutionTime(
+            SimulatedExecutionTime.ofMinimumAndRandomTime(2000, 0));
+        statement.setQueryTimeout(1);
+        assertThrows(
+            JdbcSqlTimeoutException.class,
+            () -> statement.executeUpdate(INSERT_STATEMENT.getSql()));
+      }
+    }
+  }
+
+  @Test
+  public void testExecuteBatchTimeout() throws SQLException {
+    try (java.sql.Connection connection = createJdbcConnection()) {
+      try (Statement statement = connection.createStatement()) {
+        // First verify that batch dml does not time out by default.
+        statement.addBatch(INSERT_STATEMENT.getSql());
+        int[] result = statement.executeBatch();
+        assertArrayEquals(new int[] {1}, result);
+
+        // Simulate that executeBatchDml takes 2 seconds and set a statement timeout of 1 second.
+        mockSpanner.setExecuteBatchDmlExecutionTime(
+            SimulatedExecutionTime.ofMinimumAndRandomTime(2000, 0));
+        statement.setQueryTimeout(1);
+        statement.addBatch(INSERT_STATEMENT.getSql());
+        assertThrows(JdbcSqlTimeoutException.class, statement::executeBatch);
+      }
+    }
+  }
+}