feat: support statement tags in comments in Connection API

.kokoro/nightly/integration.cfg

@@ -13,12 +13,12 @@ env_vars: {
 # TODO: remove this after we've migrated all tests and scripts
 env_vars: {
   key: "GCLOUD_PROJECT"
-  value: "cloud-java-ci-sample"
+  value: "java-docs-samples-testing"
 }
 
 env_vars: {
   key: "GOOGLE_CLOUD_PROJECT"
-  value: "cloud-java-ci-sample"
+  value: "java-docs-samples-testing"
 }
 
 env_vars: {

google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java

@@ -333,6 +333,20 @@ static final class TagOption extends InternalOption implements ReadQueryUpdateTr
     void appendToOptions(Options options) {
       options.tag = tag;
     }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(this.tag);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (!(o instanceof TagOption)) {
+        return false;
+      }
+      TagOption other = (TagOption) o;
+      return Objects.equals(this.tag, other.tag);
+    }
   }
 
   static final class EtagOption extends InternalOption implements DeleteAdminApiOption {

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

@@ -19,6 +19,8 @@
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.ErrorCode;
+import com.google.cloud.spanner.Options;
+import com.google.cloud.spanner.Options.ReadQueryUpdateTransactionOption;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.Statement;
@@ -620,6 +622,29 @@ public String removeCommentsAndTrim(String sql) {
   /** Removes any statement hints at the beginning of the statement. */
   abstract String removeStatementHint(String sql);
 
+  @VisibleForTesting
+  static final ReadQueryUpdateTransactionOption[] EMPTY_OPTIONS =
+      new ReadQueryUpdateTransactionOption[0];
+
+  /**
+   * Extracts any query/update options from the SQL string. Currently, this only supports extracting
+   * a statement tag, and the statement tag must be given as a statement hint in a comment at the
+   * start of the query string.
+   */
+  ReadQueryUpdateTransactionOption[] extractOptions(ParsedStatement statement) {
+    final String statementTagPrefix = "/*@{STATEMENT_TAG=";
+
+    String sql = statement.getStatement().getSql();
+    if (sql.startsWith(statementTagPrefix)) {
+      int endIndex = sql.indexOf("}*/", statementTagPrefix.length());
+      if (endIndex > -1) {
+        String tag = sql.substring(statementTagPrefix.length(), endIndex);
+        return new ReadQueryUpdateTransactionOption[] {Options.tag(tag)};
+      }
+    }
+    return EMPTY_OPTIONS;
+  }
+
   /** Parameter information with positional parameters translated to named parameters. */
   @InternalApi
   public static class ParametersInfo {

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

@@ -32,6 +32,7 @@
 import com.google.cloud.spanner.Mutation;
 import com.google.cloud.spanner.Options;
 import com.google.cloud.spanner.Options.QueryOption;
+import com.google.cloud.spanner.Options.ReadQueryUpdateTransactionOption;
 import com.google.cloud.spanner.Options.RpcPriority;
 import com.google.cloud.spanner.Options.UpdateOption;
 import com.google.cloud.spanner.PartitionOptions;
@@ -1138,14 +1139,17 @@ public ResultSet partitionQuery(
           "Only queries can be partitioned. Invalid statement: " + query.getSql());
     }
 
+    QueryOption[] combinedOptions =
+        concat(getStatementParser().extractOptions(parsedStatement), options);
     UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork();
     return get(
         transaction.partitionQueryAsync(
             CallType.SYNC,
             parsedStatement,
             getEffectivePartitionOptions(partitionOptions),
             mergeDataBoost(
-                mergeQueryRequestOptions(parsedStatement, mergeQueryStatementTag(options)))));
+                mergeQueryRequestOptions(
+                    parsedStatement, mergeQueryStatementTag(combinedOptions)))));
   }
 
   private PartitionOptions getEffectivePartitionOptions(
@@ -1439,6 +1443,34 @@ private List<ParsedStatement> parseUpdateStatements(Iterable<Statement> updates)
     return parsedStatements;
   }
 
+  private UpdateOption[] concat(
+      ReadQueryUpdateTransactionOption[] statementOptions, UpdateOption[] argumentOptions) {
+    if (statementOptions == null || statementOptions.length == 0) {
+      return argumentOptions;
+    }
+    if (argumentOptions == null || argumentOptions.length == 0) {
+      return statementOptions;
+    }
+    UpdateOption[] result =
+        Arrays.copyOf(statementOptions, statementOptions.length + argumentOptions.length);
+    System.arraycopy(argumentOptions, 0, result, statementOptions.length, argumentOptions.length);
+    return result;
+  }
+
+  private QueryOption[] concat(
+      ReadQueryUpdateTransactionOption[] statementOptions, QueryOption[] argumentOptions) {
+    if (statementOptions == null || statementOptions.length == 0) {
+      return argumentOptions;
+    }
+    if (argumentOptions == null || argumentOptions.length == 0) {
+      return statementOptions;
+    }
+    QueryOption[] result =
+        Arrays.copyOf(statementOptions, statementOptions.length + argumentOptions.length);
+    System.arraycopy(argumentOptions, 0, result, statementOptions.length, argumentOptions.length);
+    return result;
+  }
+
   private QueryOption[] mergeDataBoost(QueryOption... options) {
     if (this.dataBoostEnabled) {
       options = appendQueryOption(options, Options.dataBoostEnabled(true));
@@ -1515,19 +1547,20 @@ private ResultSet internalExecuteQuery(
                 && (analyzeMode != AnalyzeMode.NONE || statement.hasReturningClause())),
         "Statement must either be a query or a DML mode with analyzeMode!=NONE or returning clause");
     boolean isInternalMetadataQuery = isInternalMetadataQuery(options);
+    QueryOption[] combinedOptions = concat(getStatementParser().extractOptions(statement), options);
     UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork(isInternalMetadataQuery);
     if (autoPartitionMode
         && statement.getType() == StatementType.QUERY
         && !isInternalMetadataQuery) {
       return runPartitionedQuery(
-          statement.getStatement(), PartitionOptions.getDefaultInstance(), options);
+          statement.getStatement(), PartitionOptions.getDefaultInstance(), combinedOptions);
     }
     return get(
         transaction.executeQueryAsync(
             callType,
             statement,
             analyzeMode,
-            mergeQueryRequestOptions(statement, mergeQueryStatementTag(options))));
+            mergeQueryRequestOptions(statement, mergeQueryStatementTag(combinedOptions))));
   }
 
   private AsyncResultSet internalExecuteQueryAsync(
@@ -1542,44 +1575,54 @@ private AsyncResultSet internalExecuteQueryAsync(
     ConnectionPreconditions.checkState(
         !(autoPartitionMode && statement.getType() == StatementType.QUERY),
         "Partitioned queries cannot be executed asynchronously");
-    UnitOfWork transaction =
-        getCurrentUnitOfWorkOrStartNewUnitOfWork(isInternalMetadataQuery(options));
+    boolean isInternalMetadataQuery = isInternalMetadataQuery(options);
+    QueryOption[] combinedOptions = concat(getStatementParser().extractOptions(statement), options);
+    UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork(isInternalMetadataQuery);
     return ResultSets.toAsyncResultSet(
         transaction.executeQueryAsync(
             callType,
             statement,
             analyzeMode,
-            mergeQueryRequestOptions(statement, mergeQueryStatementTag(options))),
+            mergeQueryRequestOptions(statement, mergeQueryStatementTag(combinedOptions))),
         spanner.getAsyncExecutorProvider(),
-        options);
+        combinedOptions);
   }
 
   private ApiFuture<Long> internalExecuteUpdateAsync(
-      final CallType callType, final ParsedStatement update, UpdateOption... options) {
+      final CallType callType, final ParsedStatement update, final UpdateOption... options) {
     Preconditions.checkArgument(
         update.getType() == StatementType.UPDATE, "Statement must be an update");
+    UpdateOption[] combinedOptions = concat(getStatementParser().extractOptions(update), options);
     UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork();
     return transaction.executeUpdateAsync(
-        callType, update, mergeUpdateRequestOptions(mergeUpdateStatementTag(options)));
+        callType, update, mergeUpdateRequestOptions(mergeUpdateStatementTag(combinedOptions)));
   }
 
   private ApiFuture<ResultSet> internalAnalyzeUpdateAsync(
       final CallType callType,
       final ParsedStatement update,
-      AnalyzeMode analyzeMode,
-      UpdateOption... options) {
+      final AnalyzeMode analyzeMode,
+      final UpdateOption... options) {
     Preconditions.checkArgument(
         update.getType() == StatementType.UPDATE, "Statement must be an update");
+    UpdateOption[] combinedOptions = concat(getStatementParser().extractOptions(update), options);
     UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork();
     return transaction.analyzeUpdateAsync(
-        callType, update, analyzeMode, mergeUpdateRequestOptions(mergeUpdateStatementTag(options)));
+        callType,
+        update,
+        analyzeMode,
+        mergeUpdateRequestOptions(mergeUpdateStatementTag(combinedOptions)));
   }
 
   private ApiFuture<long[]> internalExecuteBatchUpdateAsync(
-      CallType callType, List<ParsedStatement> updates, UpdateOption... options) {
+      final CallType callType, final List<ParsedStatement> updates, final UpdateOption... options) {
     UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork();
+    UpdateOption[] combinedOptions =
+        concat(
+            getStatementParser().extractOptions(updates.isEmpty() ? null : updates.get(0)),
+            options);
     return transaction.executeBatchUpdateAsync(
-        callType, updates, mergeUpdateRequestOptions(mergeUpdateStatementTag(options)));
+        callType, updates, mergeUpdateRequestOptions(mergeUpdateStatementTag(combinedOptions)));
   }
 
   private UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork() {

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

@@ -16,8 +16,10 @@
 
 package com.google.cloud.spanner.connection;
 
+import static com.google.cloud.spanner.connection.AbstractStatementParser.EMPTY_OPTIONS;
 import static com.google.common.truth.Truth.assertThat;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotSame;
@@ -28,6 +30,8 @@
 
 import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.ErrorCode;
+import com.google.cloud.spanner.Options;
+import com.google.cloud.spanner.Options.ReadQueryUpdateTransactionOption;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
@@ -1690,6 +1694,30 @@ public void testStatementCache_ParameterizedStatement() {
     assertEquals(1, stats.hitCount());
   }
 
+  @Test
+  public void testExtractOptions() {
+    assertArrayEquals(EMPTY_OPTIONS, parser.extractOptions(parser.parse(Statement.of("select 1"))));
+
+    assertArrayEquals(
+        new ReadQueryUpdateTransactionOption[] {Options.tag("foo")},
+        parser.extractOptions(parser.parse(Statement.of("/*@{STATEMENT_TAG=foo}*/ select 1"))));
+    assertArrayEquals(
+        new ReadQueryUpdateTransactionOption[] {Options.tag("tag with space")},
+        parser.extractOptions(
+            parser.parse(Statement.of("/*@{STATEMENT_TAG=tag with space}*/ select 1"))));
+    assertArrayEquals(
+        new ReadQueryUpdateTransactionOption[] {Options.tag("foo}")},
+        parser.extractOptions(parser.parse(Statement.of("/*@{STATEMENT_TAG=foo}}*/ select 1"))));
+
+    assertArrayEquals(
+        EMPTY_OPTIONS,
+        parser.extractOptions(parser.parse(Statement.of("/*@{STATEMENT_TAG=not_a_tag*/select 1"))));
+    assertArrayEquals(
+        EMPTY_OPTIONS,
+        parser.extractOptions(
+            parser.parse(Statement.of("/*@{STATEMENT_TAG=not_a_tag} */select 1"))));
+  }
+
   static void assertUnclosedLiteral(AbstractStatementParser parser, String sql) {
     SpannerException exception =
         assertThrows(