feat: add option to indicate that a statement is the last in a transaction

README.md

@@ -19,7 +19,7 @@ If you are using Maven with [BOM][libraries-bom], add this to your pom.xml file:
     <dependency>
       <groupId>com.google.cloud</groupId>
       <artifactId>libraries-bom</artifactId>
-      <version>26.53.0</version>
+      <version>26.54.0</version>
       <type>pom</type>
       <scope>import</scope>
     </dependency>

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

@@ -698,6 +698,9 @@ ExecuteSqlRequest.Builder getExecuteSqlRequestBuilder(
     if (!isReadOnly()) {
       builder.setSeqno(getSeqNo());
     }
+    if (options.hasLastStatement()) {
+      builder.setLastStatement(options.isLastStatement());
+    }
     builder.setQueryOptions(buildQueryOptions(statement.getQueryOptions()));
     builder.setRequestOptions(buildRequestOptions(options));
     return builder;
@@ -743,6 +746,9 @@ ExecuteBatchDmlRequest.Builder getExecuteBatchDmlRequestBuilder(
     if (selector != null) {
       builder.setTransaction(selector);
     }
+    if (options.hasLastStatement()) {
+      builder.setLastStatements(options.isLastStatement());
+    }
     builder.setSeqno(getSeqNo());
     builder.setRequestOptions(buildRequestOptions(options));
     return builder;

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

@@ -108,6 +108,9 @@ public interface ReadQueryUpdateTransactionOption
   /** Marker interface to mark options applicable to Update and Write operations */
   public interface UpdateTransactionOption extends UpdateOption, TransactionOption {}
 
+  /** Marker interface for options that can be used with both executeQuery and executeUpdate. */
+  public interface QueryUpdateOption extends QueryOption, UpdateOption {}
+
   /**
    * Marker interface to mark options applicable to Create, Update and Delete operations in admin
    * API.
@@ -236,6 +239,20 @@ public static DataBoostQueryOption dataBoostEnabled(Boolean dataBoostEnabled) {
     return new DataBoostQueryOption(dataBoostEnabled);
   }
 
+  /**
+   * If set to true, this option marks the end of the transaction. The transaction should be
+   * committed or aborted after this statement executes, and attempts to execute any other requests
+   * against this transaction (including reads and queries) will be rejected. Mixing mutations with
+   * statements that are marked as the last statement is not allowed.
+   *
+   * <p>For DML statements, setting this option may cause some error reporting to be deferred until
+   * commit time (e.g. validation of unique constraints). Given this, successful execution of a DML
+   * statement should not be assumed until the transaction commits.
+   */
+  public static QueryUpdateOption lastStatement() {
+    return new LastStatementUpdateOption();
+  }
+
   /**
    * Specifying this will cause the list operation to start fetching the record from this onwards.
    */
@@ -494,6 +511,7 @@ void appendToOptions(Options options) {
   private DecodeMode decodeMode;
   private RpcOrderBy orderBy;
   private RpcLockHint lockHint;
+  private Boolean lastStatement;
 
   // Construction is via factory methods below.
   private Options() {}
@@ -630,6 +648,14 @@ OrderBy orderBy() {
     return orderBy == null ? null : orderBy.proto;
   }
 
+  boolean hasLastStatement() {
+    return lastStatement != null;
+  }
+
+  Boolean isLastStatement() {
+    return lastStatement;
+  }
+
   boolean hasLockHint() {
     return lockHint != null;
   }
@@ -694,6 +720,9 @@ public String toString() {
     if (orderBy != null) {
       b.append("orderBy: ").append(orderBy).append(' ');
     }
+    if (lastStatement != null) {
+      b.append("lastStatement: ").append(lastStatement).append(' ');
+    }
     if (lockHint != null) {
       b.append("lockHint: ").append(lockHint).append(' ');
     }
@@ -737,6 +766,7 @@ public boolean equals(Object o) {
         && Objects.equals(dataBoostEnabled(), that.dataBoostEnabled())
         && Objects.equals(directedReadOptions(), that.directedReadOptions())
         && Objects.equals(orderBy(), that.orderBy())
+        && Objects.equals(isLastStatement(), that.isLastStatement())
         && Objects.equals(lockHint(), that.lockHint());
   }
 
@@ -797,6 +827,9 @@ public int hashCode() {
     if (orderBy != null) {
       result = 31 * result + orderBy.hashCode();
     }
+    if (lastStatement != null) {
+      result = 31 * result + lastStatement.hashCode();
+    }
     if (lockHint != null) {
       result = 31 * result + lockHint.hashCode();
     }
@@ -965,4 +998,24 @@ public boolean equals(Object o) {
       return Objects.equals(filter, ((FilterOption) o).filter);
     }
   }
+
+  static final class LastStatementUpdateOption extends InternalOption implements QueryUpdateOption {
+
+    LastStatementUpdateOption() {}
+
+    @Override
+    void appendToOptions(Options options) {
+      options.lastStatement = true;
+    }
+
+    @Override
+    public int hashCode() {
+      return LastStatementUpdateOption.class.hashCode();
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      return o instanceof LastStatementUpdateOption;
+    }
+  }
 }

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

@@ -265,6 +265,9 @@ static ErrorDetails extractErrorDetails(Throwable cause) {
       if (cause instanceof ApiException) {
         return ((ApiException) cause).getErrorDetails();
       }
+      if (cause instanceof SpannerException) {
+        return ((SpannerException) cause).getErrorDetails();
+      }
       prevCause = cause;
       cause = cause.getCause();
     }

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

@@ -34,6 +34,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.QueryUpdateOption;
 import com.google.cloud.spanner.Options.UpdateOption;
 import com.google.cloud.spanner.PartitionOptions;
 import com.google.cloud.spanner.ReadOnlyTransaction;
@@ -298,7 +299,8 @@ private ApiFuture<ResultSet> executeDmlReturningAsync(
                 writeTransaction.run(
                     transaction ->
                         DirectExecuteResultSet.ofResultSet(
-                            transaction.executeQuery(update.getStatement(), options)));
+                            transaction.executeQuery(
+                                update.getStatement(), appendLastStatement(options))));
             state = UnitOfWorkState.COMMITTED;
             return resultSet;
           } catch (Throwable t) {
@@ -554,11 +556,15 @@ private ApiFuture<Tuple<Long, ResultSet>> executeTransactionalUpdateAsync(
                     transaction -> {
                       if (analyzeMode == AnalyzeMode.NONE) {
                         return Tuple.of(
-                            transaction.executeUpdate(update.getStatement(), options), null);
+                            transaction.executeUpdate(
+                                update.getStatement(), appendLastStatement(options)),
+                            null);
                       }
                       ResultSet resultSet =
                           transaction.analyzeUpdateStatement(
-                              update.getStatement(), analyzeMode.getQueryAnalyzeMode(), options);
+                              update.getStatement(),
+                              analyzeMode.getQueryAnalyzeMode(),
+                              appendLastStatement(options));
                       return Tuple.of(null, resultSet);
                     });
             state = UnitOfWorkState.COMMITTED;
@@ -582,6 +588,29 @@ private ApiFuture<Tuple<Long, ResultSet>> executeTransactionalUpdateAsync(
     return transactionalResult;
   }
 
+  private static final QueryUpdateOption[] LAST_STATEMENT_OPTIONS =
+      new QueryUpdateOption[] {Options.lastStatement()};
+
+  private static UpdateOption[] appendLastStatement(UpdateOption[] options) {
+    if (options.length == 0) {
+      return LAST_STATEMENT_OPTIONS;
+    }
+    UpdateOption[] result = new UpdateOption[options.length + 1];
+    System.arraycopy(options, 0, result, 0, options.length);
+    result[result.length - 1] = LAST_STATEMENT_OPTIONS[0];
+    return result;
+  }
+
+  private static QueryOption[] appendLastStatement(QueryOption[] options) {
+    if (options.length == 0) {
+      return LAST_STATEMENT_OPTIONS;
+    }
+    QueryOption[] result = new QueryOption[options.length + 1];
+    System.arraycopy(options, 0, result, 0, options.length);
+    result[result.length - 1] = LAST_STATEMENT_OPTIONS[0];
+    return result;
+  }
+
   /**
    * Adds a callback to the given future that retries the update statement using Partitioned DML if
    * the original statement fails with a {@link TransactionMutationLimitExceededException}.
@@ -719,7 +748,8 @@ private ApiFuture<long[]> executeTransactionalBatchUpdateAsync(
                 try {
                   long[] res =
                       transaction.batchUpdate(
-                          Iterables.transform(updates, ParsedStatement::getStatement), options);
+                          Iterables.transform(updates, ParsedStatement::getStatement),
+                          appendLastStatement(options));
                   state = UnitOfWorkState.COMMITTED;
                   return res;
                 } catch (Throwable t) {

google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java

@@ -18,6 +18,7 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -266,6 +267,42 @@ public void testGetExecuteBatchDmlRequestBuilderWithPriority() {
     assertEquals(Priority.PRIORITY_LOW, request.getRequestOptions().getPriority());
   }
 
+  @Test
+  public void testExecuteSqlLastStatement() {
+    assertFalse(
+        context
+            .getExecuteSqlRequestBuilder(
+                Statement.of("insert into test (id) values (1)"),
+                QueryMode.NORMAL,
+                Options.fromUpdateOptions(),
+                false)
+            .getLastStatement());
+    assertTrue(
+        context
+            .getExecuteSqlRequestBuilder(
+                Statement.of("insert into test (id) values (1)"),
+                QueryMode.NORMAL,
+                Options.fromUpdateOptions(Options.lastStatement()),
+                false)
+            .getLastStatement());
+  }
+
+  @Test
+  public void testExecuteBatchDmlLastStatement() {
+    assertFalse(
+        context
+            .getExecuteBatchDmlRequestBuilder(
+                Collections.singleton(Statement.of("insert into test (id) values (1)")),
+                Options.fromUpdateOptions())
+            .getLastStatements());
+    assertTrue(
+        context
+            .getExecuteBatchDmlRequestBuilder(
+                Collections.singleton(Statement.of("insert into test (id) values (1)")),
+                Options.fromUpdateOptions(Options.lastStatement()))
+            .getLastStatements());
+  }
+
   public void executeSqlRequestBuilderWithRequestOptions() {
     ExecuteSqlRequest request =
         context

google-cloud-spanner/src/test/java/com/google/cloud/spanner/OptionsTest.java

@@ -789,4 +789,22 @@ public void updateOptionsExcludeTxnFromChangeStreams() {
     assertNull(option3.withExcludeTxnFromChangeStreams());
     assertThat(option3.toString()).doesNotContain("withExcludeTxnFromChangeStreams: true");
   }
+
+  @Test
+  public void testLastStatement() {
+    Options option1 = Options.fromUpdateOptions(Options.lastStatement());
+    Options option2 = Options.fromUpdateOptions(Options.lastStatement());
+    Options option3 = Options.fromUpdateOptions();
+
+    assertEquals(option1, option2);
+    assertEquals(option1.hashCode(), option2.hashCode());
+    assertNotEquals(option1, option3);
+    assertNotEquals(option1.hashCode(), option3.hashCode());
+
+    assertTrue(option1.isLastStatement());
+    assertThat(option1.toString()).contains("lastStatement: true");
+
+    assertNull(option3.isLastStatement());
+    assertThat(option3.toString()).doesNotContain("lastStatement: true");
+  }
 }