feat: savepoints

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)