diff --git a/pom.xml b/pom.xml index ec1d8ba4..4846278c 100644 --- a/pom.xml +++ b/pom.xml @@ -241,7 +241,6 @@ com.google.cloud.spanner.jdbc.it.** - com.google.cloud.spanner.jdbc.JdbcStatementTimeoutTest sponge_log diff --git a/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java b/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java index f3a1539b..3152445b 100644 --- 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 runWithStatementTimeout(JdbcFunction 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 index a9d18400..eb876c23 100644 --- 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); }