fix: disallow statement tags for commit/rollback/run

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

@@ -19,7 +19,6 @@
 import com.google.api.core.ApiFuture;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.Options.QueryOption;
-import com.google.cloud.spanner.Options.UpdateOption;
 import com.google.cloud.spanner.ReadContext;
 import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.SpannerException;
@@ -89,7 +88,7 @@ ResultSet internalExecuteQuery(
   }
 
   @Override
-  public ApiFuture<long[]> runBatchAsync(UpdateOption... options) {
+  public ApiFuture<long[]> runBatchAsync() {
     throw SpannerExceptionFactory.newSpannerException(
         ErrorCode.FAILED_PRECONDITION, "Run batch is not supported for transactions");
   }

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

@@ -32,10 +32,10 @@
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.TimestampBound;
 import com.google.cloud.spanner.connection.StatementResult.ResultType;
+import com.google.spanner.v1.ExecuteBatchDmlRequest;
 import java.util.Iterator;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
-import javax.annotation.Nullable;
 
 /**
  * Internal connection API for Google Cloud Spanner. This interface may introduce breaking changes
@@ -338,6 +338,12 @@ public interface Connection extends AutoCloseable {
    * automatically cleared after the statement is executed. Statement tags can be used both with
    * autocommit=true and autocommit=false, and can be used for partitioned DML.
    *
+   * <p>Statement tags are not allowed before COMMIT and ROLLBACK statements.
+   *
+   * <p>Statement tags are allowed before START BATCH DML statements and will be included in the
+   * {@link ExecuteBatchDmlRequest} that is sent to Spanner. Statement tags are not allowed inside a
+   * batch.
+   *
    * @param tag The statement tag to use with the next statement that will be executed on this
    *     connection.
    */

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

@@ -506,14 +506,16 @@ public void setTransactionTag(String tag) {
   @Override
   public String getStatementTag() {
     ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
-    ConnectionPreconditions.checkState(!isDdlBatchActive(), "This connection is in a DDL batch");
+    ConnectionPreconditions.checkState(
+        !isBatchActive(), "Statement tags are not allowed inside a batch");
     return statementTag;
   }
 
   @Override
   public void setStatementTag(String tag) {
     ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
-    ConnectionPreconditions.checkState(!isDdlBatchActive(), "This connection is in a DDL batch");
+    ConnectionPreconditions.checkState(
+        !isBatchActive(), "Statement tags are not allowed inside a batch");
 
     this.statementTag = tag;
   }
@@ -699,6 +701,8 @@ public ApiFuture<Void> rollbackAsync() {
   private ApiFuture<Void> endCurrentTransactionAsync(EndTransactionMethod endTransactionMethod) {
     ConnectionPreconditions.checkState(!isBatchActive(), "This connection has an active batch");
     ConnectionPreconditions.checkState(isInTransaction(), "This connection has no transaction");
+    ConnectionPreconditions.checkState(
+        statementTag == null, "Statement tags are not supported for COMMIT or ROLLBACK");
     ApiFuture<Void> res;
     try {
       if (isTransactionStarted()) {
@@ -1053,6 +1057,7 @@ UnitOfWork createNewUnitOfWork() {
               .setTransaction(currentUnitOfWork)
               .setStatementTimeout(statementTimeout)
               .withStatementExecutor(statementExecutor)
+              .setStatementTag(statementTag)
               .build();
         case DDL_BATCH:
           return DdlBatch.newBuilder()
@@ -1170,7 +1175,7 @@ public ApiFuture<long[]> runBatchAsync() {
     ConnectionPreconditions.checkState(isBatchActive(), "This connection has no active batch");
     try {
       if (this.currentUnitOfWork != null) {
-        return this.currentUnitOfWork.runBatchAsync(mergeUpdateStatementTag());
+        return this.currentUnitOfWork.runBatchAsync();
       }
       return ApiFutures.immediateFuture(new long[0]);
     } finally {

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

@@ -212,7 +212,7 @@ public ApiFuture<Void> writeAsync(Iterable<Mutation> mutations) {
       StatementParser.INSTANCE.parse(Statement.of("RUN BATCH"));
 
   @Override
-  public ApiFuture<long[]> runBatchAsync(UpdateOption... options) {
+  public ApiFuture<long[]> runBatchAsync() {
     ConnectionPreconditions.checkState(
         state == UnitOfWorkState.STARTED, "The batch is no longer active and cannot be ran");
     if (statements.isEmpty()) {

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

@@ -23,6 +23,7 @@
 import com.google.cloud.Timestamp;
 import com.google.cloud.spanner.ErrorCode;
 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.UpdateOption;
 import com.google.cloud.spanner.ResultSet;
@@ -41,11 +42,13 @@
  */
 class DmlBatch extends AbstractBaseUnitOfWork {
   private final UnitOfWork transaction;
+  private final String statementTag;
   private final List<ParsedStatement> statements = new ArrayList<>();
   private UnitOfWorkState state = UnitOfWorkState.STARTED;
 
   static class Builder extends AbstractBaseUnitOfWork.Builder<Builder, DmlBatch> {
     private UnitOfWork transaction;
+    private String statementTag;
 
     private Builder() {}
 
@@ -55,6 +58,11 @@ Builder setTransaction(UnitOfWork transaction) {
       return this;
     }
 
+    Builder setStatementTag(String tag) {
+      this.statementTag = tag;
+      return this;
+    }
+
     @Override
     DmlBatch build() {
       Preconditions.checkState(transaction != null, "No transaction specified");
@@ -69,6 +77,7 @@ static Builder newBuilder() {
   private DmlBatch(Builder builder) {
     super(builder);
     this.transaction = builder.transaction;
+    this.statementTag = builder.statementTag;
   }
 
   @Override
@@ -154,7 +163,7 @@ public ApiFuture<Void> writeAsync(Iterable<Mutation> mutations) {
   }
 
   @Override
-  public ApiFuture<long[]> runBatchAsync(UpdateOption... options) {
+  public ApiFuture<long[]> runBatchAsync() {
     ConnectionPreconditions.checkState(
         state == UnitOfWorkState.STARTED, "The batch is no longer active and cannot be ran");
     if (statements.isEmpty()) {
@@ -168,6 +177,8 @@ public ApiFuture<long[]> runBatchAsync(UpdateOption... options) {
     // executed AFTER a Future is done, which means that a user could read the state of the Batch
     // before it has been changed.
     final SettableApiFuture<long[]> res = SettableApiFuture.create();
+    UpdateOption[] options =
+        statementTag == null ? new UpdateOption[0] : new UpdateOption[] {Options.tag(statementTag)};
     ApiFuture<long[]> updateCounts = transaction.executeBatchUpdateAsync(statements, options);
     ApiFutures.addCallback(
         updateCounts,

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

@@ -472,7 +472,7 @@ public ApiFuture<Void> rollbackAsync() {
   }
 
   @Override
-  public ApiFuture<long[]> runBatchAsync(UpdateOption... options) {
+  public ApiFuture<long[]> runBatchAsync() {
     throw SpannerExceptionFactory.newSpannerException(
         ErrorCode.FAILED_PRECONDITION, "Run batch is not supported for single-use transactions");
   }

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

@@ -91,12 +91,11 @@ public boolean isActive() {
    * batch. This method will throw a {@link SpannerException} if called for a {@link
    * Type#TRANSACTION}.
    *
-   * @param options The update options to apply to the batch update statement.
    * @return an {@link ApiFuture} containing the update counts in case of a DML batch. Returns an
    *     array containing 1 for each successful statement and 0 for each failed statement or
    *     statement that was not executed in case of a DDL batch.
    */
-  ApiFuture<long[]> runBatchAsync(UpdateOption... options);
+  ApiFuture<long[]> runBatchAsync();
 
   /**
    * Clears the currently buffered statements in this unit of work and ends the batch. This method

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

@@ -17,13 +17,17 @@
 package com.google.cloud.spanner.connection;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.fail;
 
+import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.ResultSet;
+import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.Statement;
 import com.google.spanner.v1.CommitRequest;
 import com.google.spanner.v1.ExecuteBatchDmlRequest;
 import com.google.spanner.v1.ExecuteSqlRequest;
 import java.util.Arrays;
+import org.junit.After;
 import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -32,6 +36,54 @@
 @RunWith(JUnit4.class)
 public class TaggingTest extends AbstractMockServerTest {
 
+  @After
+  public void clearRequests() {
+    mockSpanner.clearRequests();
+  }
+
+  @Test
+  public void testStatementTagNotAllowedForCommit() {
+    try (Connection connection = createConnection()) {
+      connection.setStatementTag("tag-1");
+      try {
+        connection.commit();
+        fail("missing expected exception");
+      } catch (SpannerException e) {
+        assertThat(e.getErrorCode()).isEqualTo(ErrorCode.FAILED_PRECONDITION);
+      }
+    }
+  }
+
+  @Test
+  public void testStatementTagNotAllowedForRollback() {
+    try (Connection connection = createConnection()) {
+      connection.setStatementTag("tag-1");
+      try {
+        connection.rollback();
+        fail("missing expected exception");
+      } catch (SpannerException e) {
+        assertThat(e.getErrorCode()).isEqualTo(ErrorCode.FAILED_PRECONDITION);
+      }
+    }
+  }
+
+  @Test
+  public void testStatementTagNotAllowedInsideBatch() {
+    try (Connection connection = createConnection()) {
+      for (boolean autocommit : new boolean[] {true, false}) {
+        connection.setAutocommit(autocommit);
+        connection.startBatchDml();
+        try {
+          connection.setStatementTag("tag-1");
+          fail("missing expected exception");
+        } catch (SpannerException e) {
+          assertThat(e.getErrorCode()).isEqualTo(ErrorCode.FAILED_PRECONDITION);
+        }
+        connection.abortBatch();
+      }
+    }
+  }
+
   @Test
   public void testQuery_NoTags() {
     try (Connection connection = createConnection()) {
@@ -574,15 +626,15 @@ public void testBatchUpdate_TransactionTag() {
   }
 
   @Test
-  public void testRunBatch_StatementTag() {
+  public void testDmlBatch_StatementTag() {
     try (Connection connection = createConnection()) {
       for (boolean autocommit : new boolean[] {true, false}) {
         connection.setAutocommit(autocommit);
 
+        connection.setStatementTag("batch-tag");
         connection.startBatchDml();
         connection.execute(INSERT_STATEMENT);
         connection.execute(INSERT_STATEMENT);
-        connection.setStatementTag("batch-tag");
         connection.runBatch();
 
         assertThat(mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class)).isEqualTo(1);