chore: cleanup statement execution and test all execute methods of Statement and PreparedStatement (#1321)

* test: add tests for all execute methods of Statement and PreparedStatement

Adds tests for all the different variations of the execute method on
Statement and PreparedStatement. This will allow us to make internal
changes to the way these methods are handled, while keeping control over
any behavior changes that it might introduce.

* chore: share do-with-timeout code (#1324)

Refactor the execute methods a bit so they all share the same
do-with-timeout logic.

Author: olavloite

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

@@ -24,12 +24,15 @@
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.StatementResult;
 import com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType;
+import com.google.rpc.Code;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.SQLWarning;
 import java.sql.Statement;
 import java.util.Arrays;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
+import java.util.function.Supplier;
 
 /** Base class for Cloud Spanner JDBC {@link Statement}s */
 abstract class AbstractJdbcStatement extends AbstractJdbcWrapper implements Statement {
@@ -183,21 +186,74 @@ private ResultSet executeQuery(
       QueryAnalyzeMode analyzeMode,
       QueryOption... options)
       throws SQLException {
+    Options.QueryOption[] queryOptions = getQueryOptions(options);
+    return doWithStatementTimeout(
+        () -> {
+          com.google.cloud.spanner.ResultSet resultSet;
+          if (analyzeMode == null) {
+            resultSet = connection.getSpannerConnection().executeQuery(statement, queryOptions);
+          } else {
+            resultSet = connection.getSpannerConnection().analyzeQuery(statement, analyzeMode);
+          }
+          return JdbcResultSet.of(this, resultSet);
+        });
+  }
+
+  private <T> T doWithStatementTimeout(Supplier<T> runnable) throws SQLException {
+    return doWithStatementTimeout(runnable, ignore -> Boolean.TRUE);
+  }
+
+  private <T> T doWithStatementTimeout(
+      Supplier<T> runnable, Function<T, Boolean> shouldResetTimeout) throws SQLException {
     StatementTimeout originalTimeout = setTemporaryStatementTimeout();
+    T result = null;
     try {
-      com.google.cloud.spanner.ResultSet resultSet;
-      if (analyzeMode == null) {
-        resultSet =
-            connection.getSpannerConnection().executeQuery(statement, getQueryOptions(options));
-      } else {
-        resultSet = connection.getSpannerConnection().analyzeQuery(statement, analyzeMode);
-      }
-      return JdbcResultSet.of(this, resultSet);
-    } catch (SpannerException e) {
-      throw JdbcSqlExceptionFactory.of(e);
+      result = runnable.get();
+      return result;
+    } catch (SpannerException spannerException) {
+      throw JdbcSqlExceptionFactory.of(spannerException);
     } finally {
-      resetStatementTimeout(originalTimeout);
+      if (shouldResetTimeout.apply(result)) {
+        resetStatementTimeout(originalTimeout);
+      }
+    }
+  }
+
+  /**
+   * Executes a SQL statement on the connection of this {@link Statement} as an update (DML)
+   * statement.
+   *
+   * @param statement The SQL statement to execute
+   * @return the number of rows that was inserted/updated/deleted
+   * @throws SQLException if a database error occurs, or if the number of rows affected is larger
+   *     than {@link Integer#MAX_VALUE}
+   */
+  int executeUpdate(com.google.cloud.spanner.Statement statement) throws SQLException {
+    return checkedCast(executeLargeUpdate(statement));
+  }
+
+  /**
+   * Do a checked cast from long to int. Throws a {@link SQLException} with code {@link
+   * Code#OUT_OF_RANGE} if the update count is too big to fit in an int.
+   */
+  int checkedCast(long updateCount) throws SQLException {
+    if (updateCount > Integer.MAX_VALUE) {
+      throw JdbcSqlExceptionFactory.of(
+          "update count too large for executeUpdate: " + updateCount, Code.OUT_OF_RANGE);
     }
+    return (int) updateCount;
+  }
+
+  /**
+   * Executes a SQL statement on the connection of this {@link Statement} as an update (DML)
+   * statement.
+   *
+   * @param statement The SQL statement to execute
+   * @return the number of rows that was inserted/updated/deleted
+   * @throws SQLException if a database error occurs
+   */
+  long executeLargeUpdate(com.google.cloud.spanner.Statement statement) throws SQLException {
+    return doWithStatementTimeout(() -> connection.getSpannerConnection().executeUpdate(statement));
   }
 
   /**
@@ -210,24 +266,16 @@ private ResultSet executeQuery(
    * @throws SQLException if a database error occurs.
    */
   StatementResult execute(com.google.cloud.spanner.Statement statement) throws SQLException {
-    StatementTimeout originalTimeout = setTemporaryStatementTimeout();
-    boolean mustResetTimeout = false;
-    try {
-      StatementResult result = connection.getSpannerConnection().execute(statement);
-      mustResetTimeout = !resultIsSetStatementTimeout(result);
-      if (mustResetTimeout && resultIsShowStatementTimeout(result)) {
-        // it was a 'SHOW STATEMENT_TIMEOUT statement, we need to re-run to get the correct value
-        mustResetTimeout = false;
-        result = rerunShowStatementTimeout(statement, originalTimeout);
-      }
-      return result;
-    } catch (SpannerException e) {
-      throw JdbcSqlExceptionFactory.of(e);
-    } finally {
-      if (mustResetTimeout) {
-        resetStatementTimeout(originalTimeout);
-      }
+    StatementResult statementResult =
+        doWithStatementTimeout(
+            () -> connection.getSpannerConnection().execute(statement),
+            result -> !resultIsSetStatementTimeout(result));
+    if (resultIsShowStatementTimeout(statementResult)) {
+      // We can safely re-run it without first resetting the timeout to the original value, as that
+      // has already been done by the 'doWithStatementTimeout' function.
+      return rerunShowStatementTimeout(statement);
     }
+    return statementResult;
   }
 
   /**
@@ -250,18 +298,22 @@ StatementResult execute(com.google.cloud.spanner.Statement statement) throws SQL
    *     executed was a SET STATEMENT_TIMEOUT statement.
    */
   private boolean resultIsSetStatementTimeout(StatementResult result) {
-    return result.getClientSideStatementType() == ClientSideStatementType.SET_STATEMENT_TIMEOUT;
+    return result != null
+        && result.getClientSideStatementType() == ClientSideStatementType.SET_STATEMENT_TIMEOUT;
   }
 
   private boolean resultIsShowStatementTimeout(StatementResult result) {
-    return result.getClientSideStatementType() == ClientSideStatementType.SHOW_STATEMENT_TIMEOUT;
+    return result != null
+        && result.getClientSideStatementType() == ClientSideStatementType.SHOW_STATEMENT_TIMEOUT;
   }
 
-  private StatementResult rerunShowStatementTimeout(
-      com.google.cloud.spanner.Statement statement, StatementTimeout originalTimeout)
+  private StatementResult rerunShowStatementTimeout(com.google.cloud.spanner.Statement statement)
       throws SQLException {
-    resetStatementTimeout(originalTimeout);
-    return connection.getSpannerConnection().execute(statement);
+    try {
+      return connection.getSpannerConnection().execute(statement);
+    } catch (SpannerException spannerException) {
+      throw JdbcSqlExceptionFactory.of(spannerException);
+    }
   }
 
   @Override

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

@@ -79,11 +79,7 @@ public int executeUpdate(String sql) throws SQLException {
 
   private int executeUpdate(String sql, ImmutableList<String> generatedKeysColumns)
       throws SQLException {
-    long result = executeLargeUpdate(sql, generatedKeysColumns);
-    if (result > Integer.MAX_VALUE) {
-      throw JdbcSqlExceptionFactory.of("update count too large: " + result, Code.OUT_OF_RANGE);
-    }
-    return (int) result;
+    return checkedCast(executeLargeUpdate(sql, generatedKeysColumns));
   }
 
   /**
@@ -300,11 +296,7 @@ public ResultSet getResultSet() throws SQLException {
   @Override
   public int getUpdateCount() throws SQLException {
     checkClosed();
-    if (currentUpdateCount > Integer.MAX_VALUE) {
-      throw JdbcSqlExceptionFactory.of(
-          "update count too large: " + currentUpdateCount, Code.OUT_OF_RANGE);
-    }
-    return (int) currentUpdateCount;
+    return checkedCast(currentUpdateCount);
   }
 
   /**
@@ -481,12 +473,7 @@ private long[] executeBatch(boolean large) throws SQLException {
   int[] convertUpdateCounts(long[] updateCounts) throws SQLException {
     int[] res = new int[updateCounts.length];
     for (int index = 0; index < updateCounts.length; index++) {
-      if (updateCounts[index] > Integer.MAX_VALUE) {
-        throw JdbcSqlExceptionFactory.of(
-            String.format("Update count too large for int: %d", updateCounts[index]),
-            Code.OUT_OF_RANGE);
-      }
-      res[index] = (int) updateCounts[index];
+      res[index] = checkedCast(updateCounts[index]);
     }
     return res;
   }