chore: store query timeout as a java.time.Duration (#1865)

Store the query timeout internally as a java.time.Duration, so we can use a smaller unit for testing. This again allows us to enable the timeout tests by default.
googleapis · Dec 19, 2024 · 8f587cf · 8f587cf
1 parent dee17e2
commit 8f587cf
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 15 deletions.
diff --git a/pom.xml b/pom.xml
@@ -241,7 +241,6 @@
         <configuration>
           <excludes>
             <exclude>com.google.cloud.spanner.jdbc.it.**</exclude>
-            <exclude>com.google.cloud.spanner.jdbc.JdbcStatementTimeoutTest</exclude>
           </excludes>
           <reportNameSuffix>sponge_log</reportNameSuffix>
           <systemPropertyVariables>

diff --git a/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java b/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java
@@ -24,6 +24,7 @@
 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.common.annotations.VisibleForTesting;
 import com.google.common.base.Stopwatch;
 import com.google.rpc.Code;
 import java.sql.ResultSet;
@@ -35,6 +36,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 import java.util.function.Supplier;
+import javax.annotation.Nonnull;
 
 /** Base class for Cloud Spanner JDBC {@link Statement}s */
 abstract class AbstractJdbcStatement extends AbstractJdbcWrapper implements Statement {
@@ -45,7 +47,7 @@ abstract class AbstractJdbcStatement extends AbstractJdbcWrapper implements Stat
   private boolean closeOnCompletion;
   private boolean poolable;
   private final JdbcConnection connection;
-  private int queryTimeout;
+  private Duration queryTimeout = Duration.ZERO;
 
   AbstractJdbcStatement(JdbcConnection connection) throws SQLException {
     this.connection = connection;
@@ -148,13 +150,22 @@ protected <T> T runWithStatementTimeout(JdbcFunction<Connection, T> function)
    */
   StatementTimeout setTemporaryStatementTimeout() throws SQLException {
     StatementTimeout originalTimeout = null;
-    if (getQueryTimeout() > 0) {
+    if (!getQueryTimeoutDuration().isZero()) {
       if (connection.getSpannerConnection().hasStatementTimeout()) {
         TimeUnit unit = getAppropriateTimeUnit();
         originalTimeout =
             StatementTimeout.of(connection.getSpannerConnection().getStatementTimeout(unit), unit);
       }
-      connection.getSpannerConnection().setStatementTimeout(getQueryTimeout(), TimeUnit.SECONDS);
+      Duration queryTimeout = getQueryTimeoutDuration();
+      if (queryTimeout.getNano() > 0) {
+        connection
+            .getSpannerConnection()
+            .setStatementTimeout(queryTimeout.toMillis(), TimeUnit.MILLISECONDS);
+      } else {
+        connection
+            .getSpannerConnection()
+            .setStatementTimeout(queryTimeout.getSeconds(), TimeUnit.SECONDS);
+      }
     }
     return originalTimeout;
   }
@@ -164,7 +175,7 @@ StatementTimeout setTemporaryStatementTimeout() throws SQLException {
    * has been executed.
    */
   void resetStatementTimeout(StatementTimeout originalTimeout) throws SQLException {
-    if (getQueryTimeout() > 0) {
+    if (!getQueryTimeoutDuration().isZero()) {
       if (originalTimeout == null) {
         connection.getSpannerConnection().clearStatementTimeout();
       } else {
@@ -317,14 +328,26 @@ private StatementResult rerunShowStatementTimeout(com.google.cloud.spanner.State
 
   @Override
   public int getQueryTimeout() throws SQLException {
+    return (int) getQueryTimeoutDuration().getSeconds();
+  }
+
+  @VisibleForTesting
+  @Nonnull
+  Duration getQueryTimeoutDuration() throws SQLException {
     checkClosed();
-    return queryTimeout;
+    return this.queryTimeout;
   }
 
   @Override
   public void setQueryTimeout(int seconds) throws SQLException {
+    setQueryTimeout(Duration.ofSeconds(seconds));
+  }
+
+  @VisibleForTesting
+  void setQueryTimeout(@Nonnull Duration duration) throws SQLException {
+    JdbcPreconditions.checkArgument(!duration.isNegative(), "Timeout must be >= 0");
     checkClosed();
-    this.queryTimeout = seconds;
+    this.queryTimeout = duration;
   }
 
   @Override

diff --git a/src/test/java/com/google/cloud/spanner/jdbc/JdbcStatementTimeoutTest.java b/src/test/java/com/google/cloud/spanner/jdbc/JdbcStatementTimeoutTest.java
@@ -27,17 +27,21 @@
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.time.Duration;
+import org.junit.After;
 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.
- */
+/** Tests setting a statement timeout. */
 @RunWith(JUnit4.class)
 public class JdbcStatementTimeoutTest extends AbstractMockServerTest {
 
+  @After
+  public void resetExecutionTimes() {
+    mockSpanner.removeAllExecutionTimes();
+  }
+
   @Test
   public void testExecuteTimeout() throws SQLException {
     try (java.sql.Connection connection = createJdbcConnection()) {
@@ -50,7 +54,7 @@ public void testExecuteTimeout() throws SQLException {
         // Simulate that executeSql takes 2 seconds and set a statement timeout of 1 second.
         mockSpanner.setExecuteSqlExecutionTime(
             SimulatedExecutionTime.ofMinimumAndRandomTime(2000, 0));
-        statement.setQueryTimeout(1);
+        ((JdbcStatement) statement).setQueryTimeout(Duration.ofMillis(5L));
         assertThrows(
             JdbcSqlTimeoutException.class, () -> statement.execute(INSERT_STATEMENT.getSql()));
       }
@@ -74,7 +78,7 @@ public void testExecuteQueryTimeout() throws SQLException {
         // second.
         mockSpanner.setExecuteStreamingSqlExecutionTime(
             SimulatedExecutionTime.ofMinimumAndRandomTime(2000, 0));
-        statement.setQueryTimeout(1);
+        ((JdbcStatement) statement).setQueryTimeout(Duration.ofMillis(5L));
         assertThrows(
             JdbcSqlTimeoutException.class,
             () -> statement.executeQuery(SELECT_RANDOM_STATEMENT.getSql()));
@@ -92,7 +96,7 @@ public void testExecuteUpdateTimeout() throws SQLException {
         // Simulate that executeSql takes 2 seconds and set a statement timeout of 1 second.
         mockSpanner.setExecuteSqlExecutionTime(
             SimulatedExecutionTime.ofMinimumAndRandomTime(2000, 0));
-        statement.setQueryTimeout(1);
+        ((JdbcStatement) statement).setQueryTimeout(Duration.ofMillis(5L));
         assertThrows(
             JdbcSqlTimeoutException.class,
             () -> statement.executeUpdate(INSERT_STATEMENT.getSql()));
@@ -112,7 +116,7 @@ public void testExecuteBatchTimeout() throws SQLException {
         // Simulate that executeBatchDml takes 2 seconds and set a statement timeout of 1 second.
         mockSpanner.setExecuteBatchDmlExecutionTime(
             SimulatedExecutionTime.ofMinimumAndRandomTime(2000, 0));
-        statement.setQueryTimeout(1);
+        ((JdbcStatement) statement).setQueryTimeout(Duration.ofMillis(5L));
         statement.addBatch(INSERT_STATEMENT.getSql());
         assertThrows(JdbcSqlTimeoutException.class, statement::executeBatch);
       }