feat: support statement_timeout in connection url

Author: rayudu3745

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())

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<Duration> 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<Duration> TRANSACTION_TIMEOUT =
       create(
           "transaction_timeout",

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(

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,32 @@ 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();
+      mockSpanner.setExecuteStreamingSqlExecutionTime(
+          SimulatedExecutionTime.ofMinimumAndRandomTime(500, 0));
+      connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS);
+      try (ResultSet rs = connection.executeQuery(SELECT_RANDOM_STATEMENT)) {
+        assertNotNull(rs);
+      }
+    }
+  }
+
   @Test
   public void testTimeoutExceptionReadWriteAutocommitSlowUpdate() {
     mockSpanner.setExecuteSqlExecutionTime(