diff --git a/google-cloud-spanner/clirr-ignored-differences.xml b/google-cloud-spanner/clirr-ignored-differences.xml
index e79cac4fc78..94bb5870f81 100644
--- a/google-cloud-spanner/clirr-ignored-differences.xml
+++ b/google-cloud-spanner/clirr-ignored-differences.xml
@@ -996,4 +996,16 @@
com/google/cloud/spanner/DatabaseClient
com.google.cloud.spanner.Statement$StatementFactory getStatementFactory()
+
+
+
+ 7012
+ com/google/cloud/spanner/AsyncTransactionManager
+ com.google.cloud.spanner.AsyncTransactionManager$TransactionContextFuture beginAsync(com.google.cloud.spanner.AbortedException)
+
+
+ 7012
+ com/google/cloud/spanner/TransactionManager
+ com.google.cloud.spanner.TransactionContext begin(com.google.cloud.spanner.AbortedException)
+
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbortedException.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbortedException.java
index 3e5227888d9..74e45062113 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbortedException.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbortedException.java
@@ -17,6 +17,7 @@
package com.google.cloud.spanner;
import com.google.api.gax.rpc.ApiException;
+import com.google.protobuf.ByteString;
import javax.annotation.Nullable;
/**
@@ -32,6 +33,8 @@ public class AbortedException extends SpannerException {
*/
private static final boolean IS_RETRYABLE = false;
+ private ByteString transactionID;
+
/** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */
AbortedException(
DoNotConstructDirectly token, @Nullable String message, @Nullable Throwable cause) {
@@ -46,6 +49,9 @@ public class AbortedException extends SpannerException {
@Nullable ApiException apiException,
@Nullable XGoogSpannerRequestId reqId) {
super(token, ErrorCode.ABORTED, IS_RETRYABLE, message, cause, apiException, reqId);
+ if (cause instanceof AbortedException) {
+ this.transactionID = ((AbortedException) cause).getTransactionID();
+ }
}
/**
@@ -55,4 +61,12 @@ public class AbortedException extends SpannerException {
public boolean isEmulatorOnlySupportsOneTransactionException() {
return getMessage().endsWith("The emulator only supports one transaction at a time.");
}
+
+ void setTransactionID(ByteString transactionID) {
+ this.transactionID = transactionID;
+ }
+
+ ByteString getTransactionID() {
+ return this.transactionID;
+ }
}
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java
index c6ead432046..bb5140c4755 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java
@@ -170,6 +170,21 @@ interface AsyncTransactionFunction {
*/
TransactionContextFuture beginAsync();
+ /**
+ * Initializes a new read-write transaction that is a retry of a previously aborted transaction.
+ * This method must be called before performing any operations, and it can only be invoked once
+ * per transaction lifecycle.
+ *
+ *
This method should only be used when multiplexed sessions are enabled to create a retry for
+ * a previously aborted transaction. This method can be used instead of {@link
+ * #resetForRetryAsync()} to create a retry. Using this method or {@link #resetForRetryAsync()}
+ * will have the same effect. You must pass in the {@link AbortedException} from the previous
+ * attempt to preserve the transaction's priority.
+ *
+ *
For regular sessions, this behaves the same as {@link #beginAsync()}.
+ */
+ TransactionContextFuture beginAsync(AbortedException exception);
+
/**
* Rolls back the currently active transaction. In most cases there should be no need to call this
* explicitly since {@link #close()} would automatically roll back any active transaction.
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java
index 1578de87cdb..1f7fd2a0cbe 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java
@@ -76,14 +76,27 @@ public ApiFuture closeAsync() {
@Override
public TransactionContextFutureImpl beginAsync() {
Preconditions.checkState(txn == null, "begin can only be called once");
- return new TransactionContextFutureImpl(this, internalBeginAsync(true));
+ return new TransactionContextFutureImpl(this, internalBeginAsync(true, ByteString.EMPTY));
}
- private ApiFuture internalBeginAsync(boolean firstAttempt) {
+ @Override
+ public TransactionContextFutureImpl beginAsync(AbortedException exception) {
+ Preconditions.checkState(txn == null, "begin can only be called once");
+ Preconditions.checkNotNull(exception, "AbortedException from the previous attempt is required");
+ ByteString abortedTransactionId =
+ exception.getTransactionID() != null ? exception.getTransactionID() : ByteString.EMPTY;
+ return new TransactionContextFutureImpl(this, internalBeginAsync(true, abortedTransactionId));
+ }
+
+ private ApiFuture internalBeginAsync(
+ boolean firstAttempt, ByteString abortedTransactionID) {
txnState = TransactionState.STARTED;
// Determine the latest transactionId when using a multiplexed session.
ByteString multiplexedSessionPreviousTransactionId = ByteString.EMPTY;
+ if (firstAttempt && session.getIsMultiplexed()) {
+ multiplexedSessionPreviousTransactionId = abortedTransactionID;
+ }
if (txn != null && session.getIsMultiplexed() && !firstAttempt) {
// Use the current transactionId if available, otherwise fallback to the previous aborted
// transactionId.
@@ -187,7 +200,7 @@ public TransactionContextFuture resetForRetryAsync() {
throw new IllegalStateException(
"resetForRetry can only be called if the previous attempt aborted");
}
- return new TransactionContextFutureImpl(this, internalBeginAsync(false));
+ return new TransactionContextFutureImpl(this, internalBeginAsync(false, ByteString.EMPTY));
}
@Override
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedAsyncTransactionManager.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedAsyncTransactionManager.java
index 56b874e4a87..530670960ca 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedAsyncTransactionManager.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedAsyncTransactionManager.java
@@ -50,6 +50,11 @@ public TransactionContextFuture beginAsync() {
return getAsyncTransactionManager().beginAsync();
}
+ @Override
+ public TransactionContextFuture beginAsync(AbortedException exception) {
+ return getAsyncTransactionManager().beginAsync(exception);
+ }
+
@Override
public ApiFuture rollbackAsync() {
return getAsyncTransactionManager().rollbackAsync();
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedTransactionManager.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedTransactionManager.java
index 29eae6477fc..96400e9e9bb 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedTransactionManager.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedTransactionManager.java
@@ -49,6 +49,11 @@ public TransactionContext begin() {
return getTransactionManager().begin();
}
+ @Override
+ public TransactionContext begin(AbortedException exception) {
+ return getTransactionManager().begin(exception);
+ }
+
@Override
public void commit() {
getTransactionManager().commit();
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java
index 37fa2c5d202..a7548758b35 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java
@@ -900,6 +900,13 @@ public TransactionContext begin() {
return internalBegin();
}
+ @Override
+ public TransactionContext begin(AbortedException exception) {
+ // For regular sessions, the input exception is ignored and the behavior is equivalent to
+ // calling {@link #begin()}.
+ return begin();
+ }
+
private TransactionContext internalBegin() {
TransactionContext res = new SessionPoolTransactionContext(this, delegate.begin());
session.get().markUsed();
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java
index 3d6a015531f..5e48d1b78bc 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java
@@ -163,6 +163,13 @@ public void onSuccess(TransactionContext result) {
return new TransactionContextFutureImpl(this, delegateTxnFuture);
}
+ @Override
+ public TransactionContextFuture beginAsync(AbortedException exception) {
+ // For regular sessions, the input exception is ignored and the behavior is equivalent to
+ // calling {@link #beginAsync()}.
+ return beginAsync();
+ }
+
@Override
public void onError(Throwable t) {
if (t instanceof AbortedException) {
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManager.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManager.java
index 76656efea28..350adb2a2c2 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManager.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManager.java
@@ -61,6 +61,21 @@ enum TransactionState {
*/
TransactionContext begin();
+ /**
+ * Initializes a new read-write transaction that is a retry of a previously aborted transaction.
+ * This method must be called before performing any operations, and it can only be invoked once
+ * per transaction lifecycle.
+ *
+ * This method should only be used when multiplexed sessions are enabled to create a retry for
+ * a previously aborted transaction. This method can be used instead of {@link #resetForRetry()}
+ * to create a retry. Using this method or {@link #resetForRetry()} will have the same effect. You
+ * must pass in the {@link AbortedException} from the previous attempt to preserve the
+ * transaction's priority.
+ *
+ *
For regular sessions, this behaves the same as {@link #begin()}.
+ */
+ TransactionContext begin(AbortedException exception);
+
/**
* Commits the currently active transaction. If the transaction was already aborted, then this
* would throw an {@link AbortedException}.
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java
index bbf34ab5c8f..aaed30e7fa7 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java
@@ -53,8 +53,21 @@ public void setSpan(ISpan span) {
@Override
public TransactionContext begin() {
Preconditions.checkState(txn == null, "begin can only be called once");
+ return begin(ByteString.EMPTY);
+ }
+
+ @Override
+ public TransactionContext begin(AbortedException exception) {
+ Preconditions.checkState(txn == null, "begin can only be called once");
+ Preconditions.checkNotNull(exception, "AbortedException from the previous attempt is required");
+ ByteString previousAbortedTransactionID =
+ exception.getTransactionID() != null ? exception.getTransactionID() : ByteString.EMPTY;
+ return begin(previousAbortedTransactionID);
+ }
+
+ TransactionContext begin(ByteString previousTransactionId) {
try (IScope s = tracer.withSpan(span)) {
- txn = session.newTransaction(options, /* previousTransactionId = */ ByteString.EMPTY);
+ txn = session.newTransaction(options, previousTransactionId);
session.setActive(this);
txnState = TransactionState.STARTED;
return txn;
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java
index a715fae0fad..388301d9bf7 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java
@@ -793,6 +793,11 @@ public SpannerException onError(SpannerException e, boolean withBeginTransaction
long delay = -1L;
if (exceptionToThrow instanceof AbortedException) {
delay = exceptionToThrow.getRetryDelayInMillis();
+ ((AbortedException) exceptionToThrow)
+ .setTransactionID(
+ this.transactionId != null
+ ? this.transactionId
+ : this.getPreviousTransactionId());
}
if (delay == -1L) {
txnLogger.log(
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientMockServerTest.java
index 126a2df5f42..d0841e3ac40 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientMockServerTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientMockServerTest.java
@@ -2170,6 +2170,279 @@ public void testBatchWriteAtLeastOnce() {
assertFalse(mockSpanner.getSession(executeSqlRequests.get(2).getSession()).getMultiplexed());
}
+ @Test
+ public void
+ testRWTransactionWithTransactionManager_CommitAborted_SetsTransactionId_AndUsedInNewInstance() {
+ // The below test verifies the behaviour of begin(AbortedException) method which is used to
+ // maintain transaction priority if resetForRetry() is not called.
+
+ // This test performs the following steps:
+ // 1. Simulates an ABORTED exception during commit and verifies that the transaction ID is
+ // included in the AbortedException.
+ // 2. Passes the ABORTED exception to the begin(AbortedException) method of a new
+ // TransactionManager, and verifies that the transaction ID from the failed transaction is sent
+ // during the inline begin of the first request.
+ DatabaseClientImpl client =
+ (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+ // Force the Commit RPC to return Aborted the first time it is called. The exception is cleared
+ // after the first call, so the retry should succeed.
+ mockSpanner.setCommitExecutionTime(
+ SimulatedExecutionTime.ofException(
+ mockSpanner.createAbortedException(ByteString.copyFromUtf8("test"))));
+
+ ByteString abortedTransactionID = null;
+ AbortedException exception = null;
+ try (TransactionManager manager = client.transactionManager()) {
+ TransactionContext transaction = manager.begin();
+ try {
+ try (ResultSet resultSet = transaction.executeQuery(STATEMENT)) {
+ //noinspection StatementWithEmptyBody
+ while (resultSet.next()) {
+ // ignore
+ }
+ }
+ manager.commit();
+ assertNotNull(manager.getCommitTimestamp());
+ } catch (AbortedException e) {
+ // The transactionID of the Aborted transaction should be set in AbortedException class.
+ assertNotNull(e.getTransactionID());
+ abortedTransactionID = e.getTransactionID();
+ exception = e;
+ }
+ }
+ // Verify that the transactionID of the aborted transaction is set.
+ assertNotNull(abortedTransactionID);
+ assertNotNull(exception);
+ mockSpanner.clearRequests();
+
+ // Pass AbortedException while invoking begin on the new manager instance.
+ try (TransactionManager manager = client.transactionManager()) {
+ TransactionContext transaction = manager.begin(exception);
+ while (true) {
+ try {
+ try (ResultSet resultSet = transaction.executeQuery(STATEMENT)) {
+ //noinspection StatementWithEmptyBody
+ while (resultSet.next()) {
+ // ignore
+ }
+ }
+ manager.commit();
+ assertNotNull(manager.getCommitTimestamp());
+ break;
+ } catch (AbortedException e) {
+ transaction = manager.resetForRetry();
+ }
+ }
+ }
+
+ // Verify that the ExecuteSqlRequest with the inline begin passes the transactionID of the
+ // previously aborted transaction.
+ List executeSqlRequests =
+ mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+ assertEquals(1, executeSqlRequests.size());
+ assertTrue(mockSpanner.getSession(executeSqlRequests.get(0).getSession()).getMultiplexed());
+ assertNotNull(
+ executeSqlRequests
+ .get(0)
+ .getTransaction()
+ .getBegin()
+ .getReadWrite()
+ .getMultiplexedSessionPreviousTransactionId());
+ assertEquals(
+ executeSqlRequests
+ .get(0)
+ .getTransaction()
+ .getBegin()
+ .getReadWrite()
+ .getMultiplexedSessionPreviousTransactionId(),
+ abortedTransactionID);
+
+ assertNotNull(client.multiplexedSessionDatabaseClient);
+ assertEquals(2L, client.multiplexedSessionDatabaseClient.getNumSessionsAcquired().get());
+ assertEquals(2L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get());
+ }
+
+ @Test
+ public void
+ testRWTransactionWithTransactionManager_ExecuteSQLAborted_SetsTransactionId_AndUsedInNewInstance() {
+ // This test performs the following steps:
+ // 1. Simulates an ABORTED exception during ExecuteSQL and verifies that the transaction ID is
+ // included in the AbortedException.
+ // 2. Passes the ABORTED exception to the begin(AbortedException) method of a new
+ // TransactionManager, and verifies that the transaction ID from the failed transaction is sent
+ // during the inline begin of the first request.
+ DatabaseClientImpl client =
+ (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+
+ ByteString abortedTransactionID = null;
+ AbortedException exception = null;
+ try (TransactionManager manager = client.transactionManager()) {
+ TransactionContext transaction = manager.begin();
+ try {
+ try (ResultSet resultSet = transaction.executeQuery(STATEMENT)) {
+ //noinspection StatementWithEmptyBody
+ while (resultSet.next()) {
+ // ignore
+ }
+ }
+
+ // Simulate an ABORTED in next ExecuteSQL request.
+ mockSpanner.setExecuteStreamingSqlExecutionTime(
+ SimulatedExecutionTime.ofException(
+ mockSpanner.createAbortedException(ByteString.copyFromUtf8("test"))));
+
+ try (ResultSet resultSet = transaction.executeQuery(STATEMENT)) {
+ //noinspection StatementWithEmptyBody
+ while (resultSet.next()) {
+ // ignore
+ }
+ }
+ manager.commit();
+ assertNotNull(manager.getCommitTimestamp());
+ } catch (AbortedException e) {
+ // The transactionID of the Aborted transaction should be set in AbortedException class.
+ assertNotNull(e.getTransactionID());
+ abortedTransactionID = e.getTransactionID();
+ exception = e;
+ }
+ }
+ // Verify that the transactionID of the aborted transaction is set.
+ assertNotNull(abortedTransactionID);
+ assertNotNull(exception);
+ mockSpanner.clearRequests();
+
+ // Pass AbortedException while invoking begin on the new manager instance.
+ try (TransactionManager manager = client.transactionManager()) {
+ TransactionContext transaction = manager.begin(exception);
+ while (true) {
+ try {
+ try (ResultSet resultSet = transaction.executeQuery(STATEMENT)) {
+ //noinspection StatementWithEmptyBody
+ while (resultSet.next()) {
+ // ignore
+ }
+ }
+ manager.commit();
+ assertNotNull(manager.getCommitTimestamp());
+ break;
+ } catch (AbortedException e) {
+ transaction = manager.resetForRetry();
+ }
+ }
+ }
+
+ // Verify that the ExecuteSqlRequest with inline begin includes the transaction ID from the
+ // previously aborted transaction.
+ List executeSqlRequests =
+ mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+ assertEquals(1, executeSqlRequests.size());
+ assertTrue(mockSpanner.getSession(executeSqlRequests.get(0).getSession()).getMultiplexed());
+ assertNotNull(
+ executeSqlRequests
+ .get(0)
+ .getTransaction()
+ .getBegin()
+ .getReadWrite()
+ .getMultiplexedSessionPreviousTransactionId());
+ assertEquals(
+ executeSqlRequests
+ .get(0)
+ .getTransaction()
+ .getBegin()
+ .getReadWrite()
+ .getMultiplexedSessionPreviousTransactionId(),
+ abortedTransactionID);
+
+ assertNotNull(client.multiplexedSessionDatabaseClient);
+ assertEquals(2L, client.multiplexedSessionDatabaseClient.getNumSessionsAcquired().get());
+ assertEquals(2L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get());
+ }
+
+ @Test
+ public void
+ testRWTransactionWithAsyncTransactionManager_CommitAborted_SetsTransactionId_AndUsedInNewInstance()
+ throws Exception {
+ // This test performs the following steps:
+ // 1. Simulates an ABORTED exception during ExecuteSQL and verifies that the transaction ID is
+ // included in the AbortedException.
+ // 2. Passes the ABORTED exception to the begin(AbortedException) method of a new
+ // AsyncTransactionManager, and verifies that the transaction ID from the failed transaction is
+ // sent
+ // during the inline begin of the first request.
+ DatabaseClientImpl client =
+ (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+ // Force the Commit RPC to return Aborted the first time it is called. The exception is cleared
+ // after the first call, so the retry should succeed.
+ mockSpanner.setCommitExecutionTime(
+ SimulatedExecutionTime.ofException(
+ mockSpanner.createAbortedException(ByteString.copyFromUtf8("test"))));
+ ByteString abortedTransactionID = null;
+ AbortedException exception = null;
+ try (AsyncTransactionManager manager = client.transactionManagerAsync()) {
+ TransactionContextFuture transactionContextFuture = manager.beginAsync();
+ try {
+ AsyncTransactionStep updateCount =
+ transactionContextFuture.then(
+ (transaction, ignored) -> transaction.executeUpdateAsync(UPDATE_STATEMENT),
+ MoreExecutors.directExecutor());
+ CommitTimestampFuture commitTimestamp = updateCount.commitAsync();
+ assertEquals(UPDATE_COUNT, updateCount.get().longValue());
+ assertNotNull(commitTimestamp.get());
+ } catch (AbortedException e) {
+ assertNotNull(e.getTransactionID());
+ exception = e;
+ abortedTransactionID = e.getTransactionID();
+ }
+ }
+
+ // Verify that the transactionID of the aborted transaction is set.
+ assertNotNull(abortedTransactionID);
+ assertNotNull(exception);
+ mockSpanner.clearRequests();
+
+ try (AsyncTransactionManager manager = client.transactionManagerAsync()) {
+ TransactionContextFuture transactionContextFuture = manager.beginAsync(exception);
+ while (true) {
+ try {
+ AsyncTransactionStep updateCount =
+ transactionContextFuture.then(
+ (transaction, ignored) -> transaction.executeUpdateAsync(UPDATE_STATEMENT),
+ MoreExecutors.directExecutor());
+ CommitTimestampFuture commitTimestamp = updateCount.commitAsync();
+ assertEquals(UPDATE_COUNT, updateCount.get().longValue());
+ assertNotNull(commitTimestamp.get());
+ break;
+ } catch (AbortedException e) {
+ transactionContextFuture = manager.resetForRetryAsync();
+ }
+ }
+ }
+
+ List executeSqlRequests =
+ mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+ assertEquals(1, executeSqlRequests.size());
+ assertTrue(mockSpanner.getSession(executeSqlRequests.get(0).getSession()).getMultiplexed());
+ assertNotNull(
+ executeSqlRequests
+ .get(0)
+ .getTransaction()
+ .getBegin()
+ .getReadWrite()
+ .getMultiplexedSessionPreviousTransactionId());
+ assertEquals(
+ executeSqlRequests
+ .get(0)
+ .getTransaction()
+ .getBegin()
+ .getReadWrite()
+ .getMultiplexedSessionPreviousTransactionId(),
+ abortedTransactionID);
+
+ assertNotNull(client.multiplexedSessionDatabaseClient);
+ assertEquals(2L, client.multiplexedSessionDatabaseClient.getNumSessionsAcquired().get());
+ assertEquals(2L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get());
+ }
+
private void waitForSessionToBeReplaced(DatabaseClientImpl client) {
assertNotNull(client.multiplexedSessionDatabaseClient);
SessionReference sessionReference =
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java
index d4b7c035658..ead2fd0f655 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java
@@ -46,6 +46,7 @@
import com.google.api.gax.longrunning.OperationFuture;
import com.google.cloud.NoCredentials;
import com.google.cloud.Timestamp;
+import com.google.cloud.spanner.AbortedException;
import com.google.cloud.spanner.BatchClient;
import com.google.cloud.spanner.BatchReadOnlyTransaction;
import com.google.cloud.spanner.BatchTransactionId;
@@ -120,6 +121,11 @@ public TransactionContext begin() {
return txContext;
}
+ @Override
+ public TransactionContext begin(AbortedException exception) {
+ return begin();
+ }
+
@Override
public void commit() {
Timestamp commitTimestamp = Timestamp.now();
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java
index 17994b34682..947fe9d33cf 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java
@@ -97,6 +97,11 @@ public TransactionContext begin() {
return txContext;
}
+ @Override
+ public TransactionContext begin(AbortedException exception) {
+ return begin();
+ }
+
@Override
public void commit() {
switch (commitBehavior) {
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java
index 76ace88d77a..bf4d8655d09 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java
@@ -36,6 +36,7 @@
import com.google.api.core.ApiFuture;
import com.google.api.gax.longrunning.OperationFuture;
import com.google.cloud.Timestamp;
+import com.google.cloud.spanner.AbortedException;
import com.google.cloud.spanner.AsyncResultSet;
import com.google.cloud.spanner.BatchClient;
import com.google.cloud.spanner.CommitResponse;
@@ -130,6 +131,11 @@ public TransactionContext begin() {
return txContext;
}
+ @Override
+ public TransactionContext begin(AbortedException exception) {
+ return begin();
+ }
+
@Override
public void commit() {
switch (commitBehavior) {