chore: support RESET ALL in the Connection API (#3142)

* chore: support RESET ALL in the Connection API

Adds support for the RESET ALL statement to the Connection
API. This allows users to easily reset a connection to its
original state, for example before returning a connection
to a pool, or after retrieving a connection from a pool.

Fixes googleapis/java-spanner-jdbc#1116

* 🦉 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>

Author: olavloite

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

@@ -725,4 +725,11 @@
     <className>com/google/cloud/spanner/SessionPoolOptions$Builder</className>
     <method>com.google.cloud.spanner.SessionPoolOptions$Builder setUseMultiplexedSession(boolean)</method>
   </difference>
+  
+  <!-- Added reset() method -->
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/connection/Connection</className>
+    <method>void reset()</method>
+  </difference>
 </differences>

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

@@ -172,6 +172,19 @@ public interface Connection extends AutoCloseable {
   /** @return <code>true</code> if this connection has been closed. */
   boolean isClosed();
 
+  /**
+   * Resets the state of this connection to the default state that it had when it was first created.
+   * Calling this method after a transaction has started (that is; after a statement has been
+   * executed in the transaction), does not change the active transaction. If for example a
+   * transaction has been started with a transaction tag, the transaction tag for the active
+   * transaction is not reset.
+   *
+   * <p>You can use this method to reset the state of the connection before returning a connection
+   * to a connection pool, and/or before using a connection that was retrieved from a connection
+   * pool.
+   */
+  void reset();
+
   /**
    * Sets autocommit on/off for this {@link Connection}. Connections in autocommit mode will apply
    * any changes to the database directly without waiting for an explicit commit. DDL- and DML

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

@@ -306,21 +306,10 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
     }
     this.dbClient = spanner.getDatabaseClient(options.getDatabaseId());
     this.batchClient = spanner.getBatchClient(options.getDatabaseId());
-    this.retryAbortsInternally = options.isRetryAbortsInternally();
-    this.readOnly = options.isReadOnly();
-    this.autocommit = options.isAutocommit();
-    this.queryOptions = this.queryOptions.toBuilder().mergeFrom(options.getQueryOptions()).build();
-    this.rpcPriority = options.getRPCPriority();
-    this.ddlInTransactionMode = options.getDdlInTransactionMode();
-    this.returnCommitStats = options.isReturnCommitStats();
-    this.delayTransactionStartUntilFirstWrite = options.isDelayTransactionStartUntilFirstWrite();
-    this.dataBoostEnabled = options.isDataBoostEnabled();
-    this.autoPartitionMode = options.isAutoPartitionMode();
-    this.maxPartitions = options.getMaxPartitions();
-    this.maxPartitionedParallelism = options.getMaxPartitionedParallelism();
-    this.maxCommitDelay = options.getMaxCommitDelay();
     this.ddlClient = createDdlClient();
-    setDefaultTransactionOptions();
+
+    // (Re)set the state of the connection to the default.
+    reset();
   }
 
   /** Constructor only for test purposes. */
@@ -434,6 +423,42 @@ public ApiFuture<Void> closeAsync() {
     return ApiFutures.immediateFuture(null);
   }
 
+  /**
+   * Resets the state of this connection to the default state in the {@link ConnectionOptions} of
+   * this connection.
+   */
+  public void reset() {
+    ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+
+    this.retryAbortsInternally = options.isRetryAbortsInternally();
+    this.readOnly = options.isReadOnly();
+    this.autocommit = options.isAutocommit();
+    this.queryOptions =
+        QueryOptions.getDefaultInstance().toBuilder().mergeFrom(options.getQueryOptions()).build();
+    this.rpcPriority = options.getRPCPriority();
+    this.ddlInTransactionMode = options.getDdlInTransactionMode();
+    this.returnCommitStats = options.isReturnCommitStats();
+    this.delayTransactionStartUntilFirstWrite = options.isDelayTransactionStartUntilFirstWrite();
+    this.dataBoostEnabled = options.isDataBoostEnabled();
+    this.autoPartitionMode = options.isAutoPartitionMode();
+    this.maxPartitions = options.getMaxPartitions();
+    this.maxPartitionedParallelism = options.getMaxPartitionedParallelism();
+    this.maxCommitDelay = options.getMaxCommitDelay();
+
+    this.autocommitDmlMode = AutocommitDmlMode.TRANSACTIONAL;
+    this.readOnlyStaleness = TimestampBound.strong();
+    this.statementTag = null;
+    this.statementTimeout = new StatementExecutor.StatementTimeout();
+    this.directedReadOptions = null;
+    this.savepointSupport = SavepointSupport.FAIL_AFTER_ROLLBACK;
+    this.protoDescriptors = null;
+    this.protoDescriptorsFilePath = null;
+
+    if (!isTransactionStarted()) {
+      setDefaultTransactionOptions();
+    }
+  }
+
   /** Get the current unit-of-work type of this connection. */
   UnitOfWorkType getUnitOfWorkType() {
     return unitOfWorkType;

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

@@ -128,6 +128,8 @@ StatementResult statementSetPgSessionCharacteristicsTransactionMode(
 
   StatementResult statementAbortBatch();
 
+  StatementResult statementResetAll();
+
   StatementResult statementSetRPCPriority(Priority priority);
 
   StatementResult statementShowRPCPriority();

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

@@ -20,6 +20,7 @@
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.ABORT_BATCH;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.BEGIN;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.COMMIT;
+import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.RESET_ALL;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.ROLLBACK;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.RUN_BATCH;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_AUTOCOMMIT;
@@ -528,6 +529,12 @@ public StatementResult statementAbortBatch() {
     return noResult(ABORT_BATCH);
   }
 
+  @Override
+  public StatementResult statementResetAll() {
+    getConnection().reset();
+    return noResult(RESET_ALL);
+  }
+
   @Override
   public StatementResult statementSetRPCPriority(Priority priority) {
     RpcPriority value = validRPCPriorityValues.get(priority);

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

@@ -90,6 +90,7 @@ enum ClientSideStatementType {
     START_BATCH_DML,
     RUN_BATCH,
     ABORT_BATCH,
+    RESET_ALL,
     SET_RPC_PRIORITY,
     SHOW_RPC_PRIORITY,
     SHOW_TRANSACTION_ISOLATION_LEVEL,

google-cloud-spanner/src/main/resources/com/google/cloud/spanner/connection/ClientSideStatements.json

@@ -279,6 +279,15 @@
 	  "exampleStatements": ["abort batch"],
 	  "examplePrerequisiteStatements": ["start batch ddl"]
 	},
+	{
+		"name": "RESET ALL",
+		"executorName": "ClientSideStatementNoParamExecutor",
+		"resultType": "NO_RESULT",
+		"statementType": "RESET_ALL",
+		"regex": "(?is)\\A\\s*(?:reset)(?:\\s+all)\\s*\\z",
+		"method": "statementResetAll",
+		"exampleStatements": ["reset all"]
+	},
 	{
 	  "name": "SET AUTOCOMMIT = TRUE|FALSE",
 	  "executorName": "ClientSideStatementSetExecutor",

google-cloud-spanner/src/main/resources/com/google/cloud/spanner/connection/PG_ClientSideStatements.json

@@ -328,6 +328,15 @@
       "exampleStatements": ["abort batch"],
       "examplePrerequisiteStatements": ["start batch ddl"]
     },
+    {
+      "name": "RESET ALL",
+      "executorName": "ClientSideStatementNoParamExecutor",
+      "resultType": "NO_RESULT",
+      "statementType": "RESET_ALL",
+      "regex": "(?is)\\A\\s*(?:reset)(?:\\s+all)\\s*\\z",
+      "method": "statementResetAll",
+      "exampleStatements": ["reset all"]
+    },
     {
       "name": "SET AUTOCOMMIT =|TO TRUE|FALSE",
       "executorName": "ClientSideStatementSetExecutor",

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

@@ -30,22 +30,31 @@
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.MockSpannerServiceImpl;
 import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
+import com.google.cloud.spanner.Options.RpcPriority;
 import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.SpannerOptions;
 import com.google.cloud.spanner.Statement;
+import com.google.cloud.spanner.TimestampBound;
 import com.google.common.collect.ImmutableList;
 import com.google.spanner.v1.BatchCreateSessionsRequest;
 import com.google.spanner.v1.CommitRequest;
+import com.google.spanner.v1.DirectedReadOptions;
+import com.google.spanner.v1.DirectedReadOptions.ExcludeReplicas;
+import com.google.spanner.v1.DirectedReadOptions.ReplicaSelection;
 import com.google.spanner.v1.ExecuteBatchDmlRequest;
 import com.google.spanner.v1.ExecuteSqlRequest;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
 import com.google.spanner.v1.RequestOptions;
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
 import javax.annotation.Nonnull;
 import org.junit.After;
 import org.junit.AfterClass;
@@ -228,6 +237,145 @@ public void testBatchUpdateAborted() {
         }
       }
     }
+
+    @Test
+    public void testReset() {
+      try (ConnectionImpl connection = (ConnectionImpl) createConnection()) {
+        assertResetBooleanProperty(
+            connection,
+            true,
+            connection::setRetryAbortsInternally,
+            connection::isRetryAbortsInternally);
+        assertResetBooleanProperty(
+            connection, false, connection::setReadOnly, connection::isReadOnly);
+        assertResetBooleanProperty(
+            connection, false, connection::setAutocommit, connection::isAutocommit);
+        assertResetBooleanProperty(
+            connection, false, connection::setReturnCommitStats, connection::isReturnCommitStats);
+        assertResetBooleanProperty(
+            connection,
+            false,
+            connection::setDelayTransactionStartUntilFirstWrite,
+            connection::isDelayTransactionStartUntilFirstWrite);
+        assertResetBooleanProperty(
+            connection, false, connection::setDataBoostEnabled, connection::isDataBoostEnabled);
+        assertResetBooleanProperty(
+            connection, false, connection::setAutoPartitionMode, connection::isAutoPartitionMode);
+        assertResetBooleanProperty(
+            connection,
+            false,
+            connection::setExcludeTxnFromChangeStreams,
+            connection::isExcludeTxnFromChangeStreams);
+
+        assertResetProperty(
+            connection, "", "1", connection::setOptimizerVersion, connection::getOptimizerVersion);
+        assertResetProperty(
+            connection,
+            null,
+            RpcPriority.LOW,
+            connection::setRPCPriority,
+            connection::getRPCPriority);
+        assertResetProperty(
+            connection,
+            DdlInTransactionMode.ALLOW_IN_EMPTY_TRANSACTION,
+            DdlInTransactionMode.AUTO_COMMIT_TRANSACTION,
+            connection::setDdlInTransactionMode,
+            connection::getDdlInTransactionMode);
+        assertResetProperty(
+            connection, 0, 4, connection::setMaxPartitions, connection::getMaxPartitions);
+        assertResetProperty(
+            connection,
+            1,
+            8,
+            connection::setMaxPartitionedParallelism,
+            connection::getMaxPartitionedParallelism);
+        assertResetProperty(
+            connection,
+            null,
+            Duration.ofMillis(20),
+            connection::setMaxCommitDelay,
+            connection::getMaxCommitDelay);
+        assertResetProperty(
+            connection,
+            TimestampBound.strong(),
+            TimestampBound.ofExactStaleness(10L, TimeUnit.SECONDS),
+            connection::setReadOnlyStaleness,
+            connection::getReadOnlyStaleness);
+        assertResetProperty(
+            connection, null, "tag", connection::setStatementTag, connection::getStatementTag);
+        assertResetProperty(
+            connection, null, "tag", connection::setTransactionTag, connection::getTransactionTag);
+        assertResetProperty(
+            connection,
+            null,
+            DirectedReadOptions.newBuilder()
+                .setExcludeReplicas(
+                    ExcludeReplicas.newBuilder()
+                        .addReplicaSelections(
+                            ReplicaSelection.newBuilder().setLocation("foo").build())
+                        .build())
+                .build(),
+            connection::setDirectedRead,
+            connection::getDirectedRead);
+        assertResetProperty(
+            connection,
+            SavepointSupport.FAIL_AFTER_ROLLBACK,
+            SavepointSupport.ENABLED,
+            connection::setSavepointSupport,
+            connection::getSavepointSupport);
+        assertResetProperty(
+            connection,
+            null,
+            "descriptor".getBytes(StandardCharsets.UTF_8),
+            connection::setProtoDescriptors,
+            connection::getProtoDescriptors);
+        assertResetProperty(
+            connection,
+            null,
+            "filename",
+            connection::setProtoDescriptorsFilePath,
+            connection::getProtoDescriptorsFilePath);
+
+        // Test the AutocommitDmlMode property that is only supported in auto-commit mode.
+        connection.rollback();
+        connection.setAutocommit(true);
+        assertResetProperty(
+            connection,
+            AutocommitDmlMode.TRANSACTIONAL,
+            AutocommitDmlMode.PARTITIONED_NON_ATOMIC,
+            connection::setAutocommitDmlMode,
+            connection::getAutocommitDmlMode);
+        connection.setAutocommit(false);
+
+        // Statement timeouts use a customer getter/setter, so we need to manually test that.
+        assertEquals(0L, connection.getStatementTimeout(TimeUnit.MILLISECONDS));
+        connection.setStatementTimeout(10L, TimeUnit.SECONDS);
+        assertEquals(10L, connection.getStatementTimeout(TimeUnit.SECONDS));
+        connection.reset();
+        assertEquals(0L, connection.getStatementTimeout(TimeUnit.MILLISECONDS));
+      }
+    }
+
+    private void assertResetBooleanProperty(
+        ConnectionImpl connection,
+        boolean defaultValue,
+        Consumer<Boolean> setter,
+        Supplier<Boolean> getter) {
+      assertResetProperty(connection, defaultValue, !defaultValue, setter, getter);
+    }
+
+    private <T> void assertResetProperty(
+        ConnectionImpl connection,
+        T defaultValue,
+        T alternativeValue,
+        Consumer<T> setter,
+        Supplier<T> getter) {
+      assertEquals(defaultValue, getter.get());
+      setter.accept(alternativeValue);
+      assertEquals(alternativeValue, getter.get());
+      connection.reset();
+      assertEquals(defaultValue, getter.get());
+    }
   }
 
   public static class ConnectionMinSessionsTest extends AbstractMockServerTest {