Skip to content

Commit 542c6aa

Browse files
authored
feat: support statement_timeout in connection url (#4103)
1 parent 3beea6a commit 542c6aa

File tree

4 files changed

+77
-3
lines changed

4 files changed

+77
-3
lines changed

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444
import static com.google.cloud.spanner.connection.ConnectionProperties.RETURN_COMMIT_STATS;
4545
import static com.google.cloud.spanner.connection.ConnectionProperties.RPC_PRIORITY;
4646
import static com.google.cloud.spanner.connection.ConnectionProperties.SAVEPOINT_SUPPORT;
47+
import static com.google.cloud.spanner.connection.ConnectionProperties.STATEMENT_TIMEOUT;
4748
import static com.google.cloud.spanner.connection.ConnectionProperties.TRACING_PREFIX;
4849
import static com.google.cloud.spanner.connection.ConnectionProperties.TRANSACTION_TIMEOUT;
4950

@@ -345,7 +346,7 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
345346
&& getDialect() == Dialect.POSTGRESQL
346347
? Type.TRANSACTIONAL
347348
: Type.NON_TRANSACTIONAL));
348-
349+
setInitialStatementTimeout(options.getInitialConnectionPropertyValue(STATEMENT_TIMEOUT));
349350
// (Re)set the state of the connection to the default.
350351
setDefaultTransactionOptions(getDefaultIsolationLevel());
351352
}
@@ -379,6 +380,7 @@ && getDialect() == Dialect.POSTGRESQL
379380
new ConnectionState(
380381
options.getInitialConnectionPropertyValues(),
381382
Suppliers.ofInstance(Type.NON_TRANSACTIONAL));
383+
setInitialStatementTimeout(options.getInitialConnectionPropertyValue(STATEMENT_TIMEOUT));
382384
setReadOnly(options.isReadOnly());
383385
setAutocommit(options.isAutocommit());
384386
setReturnCommitStats(options.isReturnCommitStats());
@@ -390,6 +392,21 @@ public Spanner getSpanner() {
390392
return this.spanner;
391393
}
392394

395+
private void setInitialStatementTimeout(Duration duration) {
396+
if (duration == null || duration.isZero()) {
397+
return;
398+
}
399+
com.google.protobuf.Duration protoDuration =
400+
com.google.protobuf.Duration.newBuilder()
401+
.setSeconds(duration.getSeconds())
402+
.setNanos(duration.getNano())
403+
.build();
404+
TimeUnit unit =
405+
ReadOnlyStalenessUtil.getAppropriateTimeUnit(
406+
new ReadOnlyStalenessUtil.DurationGetter(protoDuration));
407+
setStatementTimeout(ReadOnlyStalenessUtil.durationToUnits(protoDuration, unit), unit);
408+
}
409+
393410
private DdlClient createDdlClient() {
394411
return DdlClient.newBuilder()
395412
.setDatabaseAdminClient(spanner.getDatabaseAdminClient())

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -494,6 +494,15 @@ public class ConnectionProperties {
494494
.toArray(new ReadLockMode[0]),
495495
ReadLockModeConverter.INSTANCE,
496496
Context.USER);
497+
static final ConnectionProperty<Duration> STATEMENT_TIMEOUT =
498+
create(
499+
"statement_timeout",
500+
"Adds a timeout to all statements executed on this connection. "
501+
+ "This property is only used when a statement timeout is specified.",
502+
null,
503+
null,
504+
DurationConverter.INSTANCE,
505+
Context.USER);
497506
static final ConnectionProperty<Duration> TRANSACTION_TIMEOUT =
498507
create(
499508
"transaction_timeout",

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionOptionsTest.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -436,6 +436,8 @@ public void testBuilderSetUri() {
436436
"cloudspanner://spanner.googleapis.com/projects/test-project-123/instances/test-instance?autocommit=true;readonly=false");
437437
builder.setUri(
438438
"cloudspanner://spanner.googleapis.com/projects/test-project-123?autocommit=true;readonly=false");
439+
builder.setUri(
440+
"cloudspanner://spanner.googleapis.com/projects/test-project-123?statement_timeout='10s';transaction_timeout='60s'");
439441

440442
// set invalid uri's
441443
setInvalidUri(

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementTimeoutTest.java

Lines changed: 48 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -108,10 +108,12 @@ public static Object[] parameters() {
108108
@Parameter
109109
public StatementExecutorType statementExecutorType;
110110

111-
protected ITConnection createConnection() {
111+
protected ITConnection createConnection(String additionalUrlOptions) {
112+
String urlSuffix =
113+
";trackSessionLeaks=false" + (additionalUrlOptions == null ? "" : additionalUrlOptions);
112114
ConnectionOptions options =
113115
ConnectionOptions.newBuilder()
114-
.setUri(getBaseUrl() + ";trackSessionLeaks=false")
116+
.setUri(getBaseUrl() + urlSuffix)
115117
.setStatementExecutorType(statementExecutorType)
116118
.setConfigurator(
117119
optionsConfigurator -> {
@@ -135,6 +137,10 @@ protected ITConnection createConnection() {
135137
return createITConnection(options);
136138
}
137139

140+
protected ITConnection createConnection() {
141+
return createConnection("");
142+
}
143+
138144
@Before
139145
public void setup() {
140146
// Set up a connection and get the dialect to ensure that the auto-detect-dialect query has
@@ -169,6 +175,22 @@ public void testTimeoutExceptionReadOnlyAutocommit() {
169175
}
170176
}
171177

178+
@Test
179+
public void testUrlTimeoutExceptionReadOnlyAutocommit() {
180+
mockSpanner.setExecuteStreamingSqlExecutionTime(
181+
SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0));
182+
183+
try (Connection connection =
184+
createConnection(";statement_timeout='" + TIMEOUT_FOR_SLOW_STATEMENTS + "ms'")) {
185+
connection.setAutocommit(true);
186+
connection.setReadOnly(true);
187+
SpannerException e =
188+
assertThrows(
189+
SpannerException.class, () -> connection.executeQuery(SELECT_RANDOM_STATEMENT));
190+
assertEquals(ErrorCode.DEADLINE_EXCEEDED, e.getErrorCode());
191+
}
192+
}
193+
172194
@Test
173195
public void testTimeoutExceptionReadOnlyAutocommitMultipleStatements() {
174196
mockSpanner.setExecuteStreamingSqlExecutionTime(
@@ -277,6 +299,30 @@ public void testTimeoutExceptionReadWriteAutocommitMultipleStatements() {
277299
}
278300
}
279301

302+
@Test
303+
public void testUrlStatementTimeoutOverrideToSucceed() {
304+
mockSpanner.setExecuteStreamingSqlExecutionTime(
305+
SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0));
306+
307+
try (Connection connection =
308+
createConnection(";statement_timeout='" + TIMEOUT_FOR_SLOW_STATEMENTS + "ms'")) {
309+
connection.setAutocommit(true);
310+
for (int i = 0; i < 2; i++) {
311+
SpannerException e =
312+
assertThrows(
313+
SpannerException.class, () -> connection.executeQuery(SELECT_RANDOM_STATEMENT));
314+
assertEquals(ErrorCode.DEADLINE_EXCEEDED, e.getErrorCode());
315+
}
316+
317+
// Remove slow behavior and verify a fast query succeeds after overriding the timeout.
318+
mockSpanner.removeAllExecutionTimes();
319+
connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS);
320+
try (ResultSet rs = connection.executeQuery(SELECT_RANDOM_STATEMENT)) {
321+
assertNotNull(rs);
322+
}
323+
}
324+
}
325+
280326
@Test
281327
public void testTimeoutExceptionReadWriteAutocommitSlowUpdate() {
282328
mockSpanner.setExecuteSqlExecutionTime(

0 commit comments

Comments
 (0)