From f7815e314f78cdadeab5cbb2296523b756756fea Mon Sep 17 00:00:00 2001 From: Rayudu A L P Date: Wed, 1 Oct 2025 11:26:07 +0530 Subject: [PATCH] feat: support statement_timeout in connection url --- .../spanner/connection/ConnectionImpl.java | 19 ++++++- .../connection/ConnectionProperties.java | 9 ++++ .../connection/ConnectionOptionsTest.java | 2 + .../connection/StatementTimeoutTest.java | 50 ++++++++++++++++++- 4 files changed, 77 insertions(+), 3 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java index 8ae4c0fc5fd..07ded59d85f 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java @@ -44,6 +44,7 @@ import static com.google.cloud.spanner.connection.ConnectionProperties.RETURN_COMMIT_STATS; import static com.google.cloud.spanner.connection.ConnectionProperties.RPC_PRIORITY; import static com.google.cloud.spanner.connection.ConnectionProperties.SAVEPOINT_SUPPORT; +import static com.google.cloud.spanner.connection.ConnectionProperties.STATEMENT_TIMEOUT; import static com.google.cloud.spanner.connection.ConnectionProperties.TRACING_PREFIX; import static com.google.cloud.spanner.connection.ConnectionProperties.TRANSACTION_TIMEOUT; @@ -345,7 +346,7 @@ static UnitOfWorkType of(TransactionMode transactionMode) { && getDialect() == Dialect.POSTGRESQL ? Type.TRANSACTIONAL : Type.NON_TRANSACTIONAL)); - + setInitialStatementTimeout(options.getInitialConnectionPropertyValue(STATEMENT_TIMEOUT)); // (Re)set the state of the connection to the default. setDefaultTransactionOptions(getDefaultIsolationLevel()); } @@ -379,6 +380,7 @@ && getDialect() == Dialect.POSTGRESQL new ConnectionState( options.getInitialConnectionPropertyValues(), Suppliers.ofInstance(Type.NON_TRANSACTIONAL)); + setInitialStatementTimeout(options.getInitialConnectionPropertyValue(STATEMENT_TIMEOUT)); setReadOnly(options.isReadOnly()); setAutocommit(options.isAutocommit()); setReturnCommitStats(options.isReturnCommitStats()); @@ -390,6 +392,21 @@ public Spanner getSpanner() { return this.spanner; } + private void setInitialStatementTimeout(Duration duration) { + if (duration == null || duration.isZero()) { + return; + } + com.google.protobuf.Duration protoDuration = + com.google.protobuf.Duration.newBuilder() + .setSeconds(duration.getSeconds()) + .setNanos(duration.getNano()) + .build(); + TimeUnit unit = + ReadOnlyStalenessUtil.getAppropriateTimeUnit( + new ReadOnlyStalenessUtil.DurationGetter(protoDuration)); + setStatementTimeout(ReadOnlyStalenessUtil.durationToUnits(protoDuration, unit), unit); + } + private DdlClient createDdlClient() { return DdlClient.newBuilder() .setDatabaseAdminClient(spanner.getDatabaseAdminClient()) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java index 3bb5d71e485..ae64f44ebce 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java @@ -494,6 +494,15 @@ public class ConnectionProperties { .toArray(new ReadLockMode[0]), ReadLockModeConverter.INSTANCE, Context.USER); + static final ConnectionProperty STATEMENT_TIMEOUT = + create( + "statement_timeout", + "Adds a timeout to all statements executed on this connection. " + + "This property is only used when a statement timeout is specified.", + null, + null, + DurationConverter.INSTANCE, + Context.USER); static final ConnectionProperty TRANSACTION_TIMEOUT = create( "transaction_timeout", diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionOptionsTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionOptionsTest.java index e9af95d9de0..18fd5cb6143 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionOptionsTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionOptionsTest.java @@ -436,6 +436,8 @@ public void testBuilderSetUri() { "cloudspanner://spanner.googleapis.com/projects/test-project-123/instances/test-instance?autocommit=true;readonly=false"); builder.setUri( "cloudspanner://spanner.googleapis.com/projects/test-project-123?autocommit=true;readonly=false"); + builder.setUri( + "cloudspanner://spanner.googleapis.com/projects/test-project-123?statement_timeout='10s';transaction_timeout='60s'"); // set invalid uri's setInvalidUri( diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementTimeoutTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementTimeoutTest.java index b1adb3861bc..e854b3d9d90 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementTimeoutTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementTimeoutTest.java @@ -108,10 +108,12 @@ public static Object[] parameters() { @Parameter public StatementExecutorType statementExecutorType; - protected ITConnection createConnection() { + protected ITConnection createConnection(String additionalUrlOptions) { + String urlSuffix = + ";trackSessionLeaks=false" + (additionalUrlOptions == null ? "" : additionalUrlOptions); ConnectionOptions options = ConnectionOptions.newBuilder() - .setUri(getBaseUrl() + ";trackSessionLeaks=false") + .setUri(getBaseUrl() + urlSuffix) .setStatementExecutorType(statementExecutorType) .setConfigurator( optionsConfigurator -> { @@ -135,6 +137,10 @@ protected ITConnection createConnection() { return createITConnection(options); } + protected ITConnection createConnection() { + return createConnection(""); + } + @Before public void setup() { // Set up a connection and get the dialect to ensure that the auto-detect-dialect query has @@ -169,6 +175,22 @@ public void testTimeoutExceptionReadOnlyAutocommit() { } } + @Test + public void testUrlTimeoutExceptionReadOnlyAutocommit() { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + try (Connection connection = + createConnection(";statement_timeout='" + TIMEOUT_FOR_SLOW_STATEMENTS + "ms'")) { + connection.setAutocommit(true); + connection.setReadOnly(true); + SpannerException e = + assertThrows( + SpannerException.class, () -> connection.executeQuery(SELECT_RANDOM_STATEMENT)); + assertEquals(ErrorCode.DEADLINE_EXCEEDED, e.getErrorCode()); + } + } + @Test public void testTimeoutExceptionReadOnlyAutocommitMultipleStatements() { mockSpanner.setExecuteStreamingSqlExecutionTime( @@ -277,6 +299,30 @@ public void testTimeoutExceptionReadWriteAutocommitMultipleStatements() { } } + @Test + public void testUrlStatementTimeoutOverrideToSucceed() { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + try (Connection connection = + createConnection(";statement_timeout='" + TIMEOUT_FOR_SLOW_STATEMENTS + "ms'")) { + connection.setAutocommit(true); + for (int i = 0; i < 2; i++) { + SpannerException e = + assertThrows( + SpannerException.class, () -> connection.executeQuery(SELECT_RANDOM_STATEMENT)); + assertEquals(ErrorCode.DEADLINE_EXCEEDED, e.getErrorCode()); + } + + // Remove slow behavior and verify a fast query succeeds after overriding the timeout. + mockSpanner.removeAllExecutionTimes(); + connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); + try (ResultSet rs = connection.executeQuery(SELECT_RANDOM_STATEMENT)) { + assertNotNull(rs); + } + } + } + @Test public void testTimeoutExceptionReadWriteAutocommitSlowUpdate() { mockSpanner.setExecuteSqlExecutionTime(