feat: add option to indicate that a statement is the last in a transaction (#3644)

* feat: Add LastStatement DML option

* Removing debugging changes

---------

Co-authored-by: Shirdon Gorse <[email protected]>

Author: olavloite

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.hasLastStatementSet()) {
+      builder.setLastStatement(options.lastStatementSet());
+    }
     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.hasLastStatementSet()) {
+      builder.setLastStatements(options.lastStatementSet());
+    }
     builder.setSeqno(getSeqNo());
     builder.setRequestOptions(buildRequestOptions(options));
     return builder;

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

@@ -236,6 +236,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 LastStatementUpdateOption lastStatementSet(Boolean lastStatementSet) {
+    return new LastStatementUpdateOption(lastStatementSet);
+  }
+
   /**
    * Specifying this will cause the list operation to start fetching the record from this onwards.
    */
@@ -495,6 +509,8 @@ void appendToOptions(Options options) {
   private RpcOrderBy orderBy;
   private RpcLockHint lockHint;
 
+  private Boolean lastStatementSet;
+
   // Construction is via factory methods below.
   private Options() {}
 
@@ -630,6 +646,14 @@ OrderBy orderBy() {
     return orderBy == null ? null : orderBy.proto;
   }
 
+  boolean hasLastStatementSet() {
+    return lastStatementSet != null;
+  }
+
+  Boolean lastStatementSet() {
+    return lastStatementSet;
+  }
+
   boolean hasLockHint() {
     return lockHint != null;
   }
@@ -694,6 +718,9 @@ public String toString() {
     if (orderBy != null) {
       b.append("orderBy: ").append(orderBy).append(' ');
     }
+    if (lastStatementSet != null) {
+      b.append("lastStatementSet: ").append(lastStatementSet).append(' ');
+    }
     if (lockHint != null) {
       b.append("lockHint: ").append(lockHint).append(' ');
     }
@@ -737,6 +764,7 @@ public boolean equals(Object o) {
         && Objects.equals(dataBoostEnabled(), that.dataBoostEnabled())
         && Objects.equals(directedReadOptions(), that.directedReadOptions())
         && Objects.equals(orderBy(), that.orderBy())
+        && Objects.equals(lastStatementSet(), that.lastStatementSet());
         && Objects.equals(lockHint(), that.lockHint());
   }
 
@@ -797,6 +825,9 @@ public int hashCode() {
     if (orderBy != null) {
       result = 31 * result + orderBy.hashCode();
     }
+    if (lastStatementSet != null) {
+      result = 31 * result + lastStatementSet.hashCode();
+    }
     if (lockHint != null) {
       result = 31 * result + lockHint.hashCode();
     }
@@ -965,4 +996,18 @@ public boolean equals(Object o) {
       return Objects.equals(filter, ((FilterOption) o).filter);
     }
   }
+
+  static final class LastStatementUpdateOption extends InternalOption implements UpdateOption {
+
+    private final Boolean lastStatementSet;
+
+    LastStatementUpdateOption(Boolean lastStatementSet) {
+      this.lastStatementSet = lastStatementSet;
+    }
+
+    @Override
+    void appendToOptions(Options options) {
+      options.lastStatementSet = lastStatementSet;
+    }
+  }
 }