feat: savepoints (#2278)

* feat: savepoint

Adds support for savepoints to the Connection API. Savepoints use the internal retry
mechanism for read/write transactions to emulate actual savepoints. Rolling back to
a savepoint is guaranteed to always work, but resuming the transaction after rolling
back can fail if the underlying data has been modified, or if one or more of the
statements before the savepoint that was rolled back to returns non-deterministic
data.

* fix: add clirr diff

* docs: comments + more tests

* chore: fix and test identifier verification

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* feat: make savepoint feature configurable

* fix: set savepoint support

* docs: improve javadoc

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

---------

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

README.md

@@ -57,13 +57,13 @@ implementation 'com.google.cloud:google-cloud-spanner'
 If you are using Gradle without BOM, add this to your dependencies:
 
 ```Groovy
-implementation 'com.google.cloud:google-cloud-spanner:6.38.2'
+implementation 'com.google.cloud:google-cloud-spanner:6.39.0'
 ```
 
 If you are using SBT, add this to your dependencies:
 
 ```Scala
-libraryDependencies += "com.google.cloud" % "google-cloud-spanner" % "6.38.2"
+libraryDependencies += "com.google.cloud" % "google-cloud-spanner" % "6.39.0"
 ```
 <!-- {x-version-update-end} -->
 
@@ -411,7 +411,7 @@ Java is a registered trademark of Oracle and/or its affiliates.
 [kokoro-badge-link-5]: http://storage.googleapis.com/cloud-devrel-public/java/badges/java-spanner/java11.html
 [stability-image]: https://img.shields.io/badge/stability-stable-green
 [maven-version-image]: https://img.shields.io/maven-central/v/com.google.cloud/google-cloud-spanner.svg
-[maven-version-link]: https://central.sonatype.com/artifact/com.google.cloud/google-cloud-spanner/6.38.2
+[maven-version-link]: https://central.sonatype.com/artifact/com.google.cloud/google-cloud-spanner/6.39.0
 [authentication]: https://github.com/googleapis/google-cloud-java#authentication
 [auth-scopes]: https://developers.google.com/identity/protocols/oauth2/scopes
 [predefined-iam-roles]: https://cloud.google.com/iam/docs/understanding-roles#predefined_roles

google-cloud-spanner/clirr-ignored-differences.xml

@@ -222,4 +222,30 @@
     <className>com/google/cloud/spanner/connection/Connection</className>
     <method>com.google.cloud.spanner.ResultSet analyzeUpdateStatement(com.google.cloud.spanner.Statement, com.google.cloud.spanner.ReadContext$QueryAnalyzeMode, com.google.cloud.spanner.Options$UpdateOption[])</method>
   </difference>
+  <!-- Savepoints -->
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/connection/Connection</className>
+    <method>void setSavepointSupport(com.google.cloud.spanner.connection.SavepointSupport)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/connection/Connection</className>
+    <method>com.google.cloud.spanner.connection.SavepointSupport getSavepointSupport()</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/connection/Connection</className>
+    <method>void savepoint(java.lang.String)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/connection/Connection</className>
+    <method>void releaseSavepoint(java.lang.String)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/connection/Connection</className>
+    <method>void rollbackToSavepoint(java.lang.String)</method>
+  </difference>
 </differences>

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

@@ -21,6 +21,7 @@
 import com.google.api.gax.grpc.GrpcCallContext;
 import com.google.api.gax.longrunning.OperationFuture;
 import com.google.api.gax.rpc.ApiCallContext;
+import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.Options.RpcPriority;
 import com.google.cloud.spanner.SpannerException;
@@ -45,6 +46,7 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.GuardedBy;
 
@@ -128,6 +130,33 @@ B setRpcPriority(@Nullable RpcPriority rpcPriority) {
     this.rpcPriority = builder.rpcPriority;
   }
 
+  /**
+   * Returns a descriptive name for the type of transaction / unit of work. This is used in error
+   * messages.
+   */
+  abstract String getUnitOfWorkName();
+
+  @Override
+  public void savepoint(@Nonnull String name, @Nonnull Dialect dialect) {
+    throw SpannerExceptionFactory.newSpannerException(
+        ErrorCode.FAILED_PRECONDITION, "Savepoint is not supported for " + getUnitOfWorkName());
+  }
+
+  @Override
+  public void releaseSavepoint(@Nonnull String name) {
+    throw SpannerExceptionFactory.newSpannerException(
+        ErrorCode.FAILED_PRECONDITION,
+        "Release savepoint is not supported for " + getUnitOfWorkName());
+  }
+
+  @Override
+  public void rollbackToSavepoint(
+      @Nonnull String name, @Nonnull SavepointSupport savepointSupport) {
+    throw SpannerExceptionFactory.newSpannerException(
+        ErrorCode.FAILED_PRECONDITION,
+        "Rollback to savepoint is not supported for " + getUnitOfWorkName());
+  }
+
   StatementExecutor getStatementExecutor() {
     return statementExecutor;
   }

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

@@ -17,22 +17,71 @@
 package com.google.cloud.spanner.connection;
 
 import com.google.api.core.ApiFuture;
+import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.Options.QueryOption;
 import com.google.cloud.spanner.ReadContext;
 import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
 import com.google.spanner.v1.SpannerGrpc;
+import java.util.LinkedList;
+import java.util.Objects;
+import javax.annotation.Nonnull;
 
 /**
  * Base class for {@link Connection}-based transactions that can be used for multiple read and
  * read/write statements.
  */
 abstract class AbstractMultiUseTransaction extends AbstractBaseUnitOfWork {
 
+  /** In-memory savepoint implementation that is used by the Connection API. */
+  static class Savepoint {
+    private final String name;
+
+    static Savepoint of(String name) {
+      return new Savepoint(name);
+    }
+
+    Savepoint(String name) {
+      this.name = name;
+    }
+
+    /** Returns the index of the first statement that was executed after this savepoint. */
+    int getStatementPosition() {
+      return -1;
+    }
+
+    /** Returns the index of the first mutation that was executed after this savepoint. */
+    int getMutationPosition() {
+      return -1;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (!(o instanceof Savepoint)) {
+        return false;
+      }
+      return Objects.equals(((Savepoint) o).name, name);
+    }
+
+    @Override
+    public int hashCode() {
+      return name.hashCode();
+    }
+
+    @Override
+    public String toString() {
+      return name;
+    }
+  }
+
+  private final LinkedList<Savepoint> savepoints = new LinkedList<>();
+
   AbstractMultiUseTransaction(Builder<?, ? extends AbstractMultiUseTransaction> builder) {
     super(builder);
   }
@@ -94,4 +143,53 @@ public void abortBatch() {
     throw SpannerExceptionFactory.newSpannerException(
         ErrorCode.FAILED_PRECONDITION, "Run batch is not supported for transactions");
   }
+
+  abstract Savepoint savepoint(String name);
+
+  abstract void rollbackToSavepoint(Savepoint savepoint);
+
+  @VisibleForTesting
+  ImmutableList<Savepoint> getSavepoints() {
+    return ImmutableList.copyOf(savepoints);
+  }
+
+  @Override
+  public void savepoint(@Nonnull String name, @Nonnull Dialect dialect) {
+    if (dialect != Dialect.POSTGRESQL) {
+      // Check that there is no savepoint with this name. Note that PostgreSQL allows multiple
+      // savepoints in a transaction with the same name, so we don't execute this check for PG.
+      if (savepoints.stream().anyMatch(savepoint -> savepoint.name.equals(name))) {
+        throw SpannerExceptionFactory.newSpannerException(
+            ErrorCode.INVALID_ARGUMENT, "Savepoint with name " + name + " already exists");
+      }
+    }
+    savepoints.add(savepoint(name));
+  }
+
+  @Override
+  public void releaseSavepoint(@Nonnull String name) {
+    // Remove the given savepoint and all later savepoints from the transaction.
+    savepoints.subList(getSavepointIndex(name), savepoints.size()).clear();
+  }
+
+  @Override
+  public void rollbackToSavepoint(
+      @Nonnull String name, @Nonnull SavepointSupport savepointSupport) {
+    int index = getSavepointIndex(name);
+    rollbackToSavepoint(savepoints.get(index));
+    if (index < (savepoints.size() - 1)) {
+      // Remove all savepoints that come after this savepoint from the transaction.
+      // Rolling back to a savepoint does not remove that savepoint, only the ones that come after.
+      savepoints.subList(index + 1, savepoints.size()).clear();
+    }
+  }
+
+  private int getSavepointIndex(String name) {
+    int index = savepoints.lastIndexOf(savepoint(name));
+    if (index == -1) {
+      throw SpannerExceptionFactory.newSpannerException(
+          ErrorCode.INVALID_ARGUMENT, "Savepoint with name " + name + " does not exist");
+    }
+    return index;
+  }
 }

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

@@ -457,6 +457,25 @@ public Priority convert(String value) {
     }
   }
 
+  /** Converter for converting strings to {@link SavepointSupport} values. */
+  static class SavepointSupportConverter
+      implements ClientSideStatementValueConverter<SavepointSupport> {
+    private final CaseInsensitiveEnumMap<SavepointSupport> values =
+        new CaseInsensitiveEnumMap<>(SavepointSupport.class);
+
+    public SavepointSupportConverter(String allowedValues) {}
+
+    @Override
+    public Class<SavepointSupport> getParameterClass() {
+      return SavepointSupport.class;
+    }
+
+    @Override
+    public SavepointSupport convert(String value) {
+      return values.get(value);
+    }
+  }
+
   static class ExplainCommandConverter implements ClientSideStatementValueConverter<String> {
     @Override
     public Class<String> getParameterClass() {

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

@@ -713,6 +713,61 @@ default RpcPriority getRPCPriority() {
    */
   ApiFuture<Void> rollbackAsync();
 
+  /** Returns the current savepoint support for this connection. */
+  SavepointSupport getSavepointSupport();
+
+  /** Sets how savepoints should be supported on this connection. */
+  void setSavepointSupport(SavepointSupport savepointSupport);
+
+  /**
+   * Creates a savepoint with the given name.
+   *
+   * <p>The uniqueness constraints on a savepoint name depends on the database dialect that is used:
+   *
+   * <ul>
+   *   <li>{@link Dialect#GOOGLE_STANDARD_SQL} requires that savepoint names are unique within a
+   *       transaction. The name of a savepoint that has been released or destroyed because the
+   *       transaction has rolled back to a savepoint that was defined before that savepoint can be
+   *       re-used within the transaction.
+   *   <li>{@link Dialect#POSTGRESQL} follows the rules for savepoint names in PostgreSQL. This
+   *       means that multiple savepoints in one transaction can have the same name, but only the
+   *       last savepoint with a given name is visible. See <a
+   *       href="https://www.postgresql.org/docs/current/sql-savepoint.html">PostgreSQL savepoint
+   *       documentation</a> for more information.
+   * </ul>
+   *
+   * @param name the name of the savepoint to create
+   * @throws SpannerException if a savepoint with the same name already exists and the dialect that
+   *     is used is {@link Dialect#GOOGLE_STANDARD_SQL}
+   * @throws SpannerException if there is no transaction on this connection
+   * @throws SpannerException if internal retries have been disabled for this connection
+   */
+  void savepoint(String name);
+
+  /**
+   * Releases the savepoint with the given name. The savepoint and all later savepoints will be
+   * removed from the current transaction and can no longer be used.
+   *
+   * @param name the name of the savepoint to release
+   * @throws SpannerException if no savepoint with the given name exists
+   */
+  void releaseSavepoint(String name);
+
+  /**
+   * Rolls back to the given savepoint. Rolling back to a savepoint undoes all changes and releases
+   * all internal locks that have been taken by the transaction after the savepoint. Rolling back to
+   * a savepoint does not remove the savepoint from the transaction, and it is possible to roll back
+   * to the same savepoint multiple times. All savepoints that have been defined after the given
+   * savepoint are removed from the transaction.
+   *
+   * @param name the name of the savepoint to roll back to.
+   * @throws SpannerException if no savepoint with the given name exists.
+   * @throws AbortedDueToConcurrentModificationException if rolling back to the savepoint failed
+   *     because another transaction has modified the data that has been read or modified by this
+   *     transaction
+   */
+  void rollbackToSavepoint(String name);
+
   /**
    * @return <code>true</code> if this connection has a transaction (that has not necessarily
    *     started). This method will only return false when the {@link Connection} is in autocommit

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

@@ -17,6 +17,7 @@
 package com.google.cloud.spanner.connection;
 
 import static com.google.cloud.spanner.SpannerApiFutures.get;
+import static com.google.cloud.spanner.connection.ConnectionPreconditions.checkValidIdentifier;
 
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
@@ -213,6 +214,7 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
   private TimestampBound readOnlyStaleness = TimestampBound.strong();
   private QueryOptions queryOptions = QueryOptions.getDefaultInstance();
   private RpcPriority rpcPriority = null;
+  private SavepointSupport savepointSupport = SavepointSupport.FAIL_AFTER_ROLLBACK;
 
   private String transactionTag;
   private String statementTag;
@@ -840,6 +842,46 @@ private ApiFuture<Void> endCurrentTransactionAsync(EndTransactionMethod endTrans
     return res;
   }
 
+  @Override
+  public SavepointSupport getSavepointSupport() {
+    return this.savepointSupport;
+  }
+
+  @Override
+  public void setSavepointSupport(SavepointSupport savepointSupport) {
+    ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+    ConnectionPreconditions.checkState(
+        !isBatchActive(), "Cannot set SavepointSupport while in a batch");
+    ConnectionPreconditions.checkState(
+        !isTransactionStarted(), "Cannot set SavepointSupport while a transaction is active");
+    this.savepointSupport = savepointSupport;
+  }
+
+  @Override
+  public void savepoint(String name) {
+    ConnectionPreconditions.checkState(isInTransaction(), "This connection has no transaction");
+    ConnectionPreconditions.checkState(
+        savepointSupport.isSavepointCreationAllowed(),
+        "This connection does not allow the creation of savepoints. Current value of SavepointSupport: "
+            + savepointSupport);
+    getCurrentUnitOfWorkOrStartNewUnitOfWork().savepoint(checkValidIdentifier(name), getDialect());
+  }
+
+  @Override
+  public void releaseSavepoint(String name) {
+    ConnectionPreconditions.checkState(
+        isTransactionStarted(), "This connection has no active transaction");
+    getCurrentUnitOfWorkOrStartNewUnitOfWork().releaseSavepoint(checkValidIdentifier(name));
+  }
+
+  @Override
+  public void rollbackToSavepoint(String name) {
+    ConnectionPreconditions.checkState(
+        isTransactionStarted(), "This connection has no active transaction");
+    getCurrentUnitOfWorkOrStartNewUnitOfWork()
+        .rollbackToSavepoint(checkValidIdentifier(name), savepointSupport);
+  }
+
   @Override
   public StatementResult execute(Statement statement) {
     Preconditions.checkNotNull(statement);
@@ -1302,6 +1344,7 @@ UnitOfWork createNewUnitOfWork() {
           return ReadWriteTransaction.newBuilder()
               .setDatabaseClient(dbClient)
               .setRetryAbortsInternally(retryAbortsInternally)
+              .setSavepointSupport(savepointSupport)
               .setReturnCommitStats(returnCommitStats)
               .setTransactionRetryListeners(transactionRetryListeners)
               .setStatementTimeout(statementTimeout)