chore(spanner): fix code and tests (googleapis#3636)

Author: harshachinta

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

@@ -55,7 +55,7 @@ public void setSpan(ISpan span) {
 
  @Override
  public void close() {
- closeAsync();
+ SpannerApiFutures.get(closeAsync());
  }
 
  @Override
@@ -183,6 +183,10 @@ public ApiFuture<Void> rollbackAsync() {
 
  @Override
  public TransactionContextFuture resetForRetryAsync() {
+ if (txn == null || !txn.isAborted() && txnState != TransactionState.ABORTED) {
+ throw new IllegalStateException(
+ "resetForRetry can only be called if the previous attempt aborted");
+ }
  return new TransactionContextFutureImpl(this, internalBeginAsync(false));
  }
 

google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerImplTest.java

@@ -16,18 +16,14 @@
 
 package com.google.cloud.spanner;
 
-import static org.junit.Assert.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.clearInvocations;
-import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import com.google.api.core.ApiFutures;
 import com.google.cloud.Timestamp;
-import com.google.protobuf.ByteString;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Scope;
 import org.junit.Test;
@@ -60,67 +56,4 @@ public void testCommitReturnsCommitStats() {
  verify(transaction).commitAsync();
  }
  }
-
- @Test
- public void testRetryUsesPreviousTransactionIdOnMultiplexedSession() {
- // Set up mock transaction IDs
- final ByteString mockTransactionId = ByteString.copyFromUtf8("mockTransactionId");
- final ByteString mockPreviousTransactionId =
- ByteString.copyFromUtf8("mockPreviousTransactionId");
-
- Span oTspan = mock(Span.class);
- ISpan span = new OpenTelemetrySpan(oTspan);
- when(oTspan.makeCurrent()).thenReturn(mock(Scope.class));
- // Mark the session as multiplexed.
- when(session.getIsMultiplexed()).thenReturn(true);
-
- // Initialize a mock transaction with transactionId = null, previousTransactionId = null.
- transaction = mock(TransactionRunnerImpl.TransactionContextImpl.class);
- when(transaction.ensureTxnAsync()).thenReturn(ApiFutures.immediateFuture(null));
- when(session.newTransaction(eq(Options.fromTransactionOptions(Options.commitStats())), any()))
- .thenReturn(transaction);
-
- // Simulate an ABORTED error being thrown when `commitAsync()` is called.
- doThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.ABORTED, ""))
- .when(transaction)
- .commitAsync();
-
- try (AsyncTransactionManagerImpl manager =
- new AsyncTransactionManagerImpl(session, span, Options.commitStats())) {
- manager.beginAsync();
-
- // Verify that for the first transaction attempt, the `previousTransactionId` is
- // ByteString.EMPTY.
- // This is because no transaction has been previously aborted at this point.
- verify(session)
- .newTransaction(Options.fromTransactionOptions(Options.commitStats()), ByteString.EMPTY);
- assertThrows(AbortedException.class, manager::commitAsync);
- clearInvocations(session);
-
- // Mock the transaction object to contain transactionID=null and
- // previousTransactionId=mockPreviousTransactionId
- when(transaction.getPreviousTransactionId()).thenReturn(mockPreviousTransactionId);
- manager.resetForRetryAsync();
- // Verify that in the first retry attempt, the `previousTransactionId`
- // (mockPreviousTransactionId) is passed to the new transaction.
- // This allows Spanner to retry the transaction using the ID of the aborted transaction.
- verify(session)
- .newTransaction(
- Options.fromTransactionOptions(Options.commitStats()), mockPreviousTransactionId);
- assertThrows(AbortedException.class, manager::commitAsync);
- clearInvocations(session);
-
- // Mock the transaction object to contain transactionID=mockTransactionId and
- // previousTransactionId=mockPreviousTransactionId and transactionID = null
- transaction.transactionId = mockTransactionId;
- manager.resetForRetryAsync();
- // Verify that the latest `transactionId` (mockTransactionId) is used in the retry.
- // This ensures the retry logic is working as expected with the latest transaction ID.
- verify(session)
- .newTransaction(Options.fromTransactionOptions(Options.commitStats()), mockTransactionId);
-
- when(transaction.rollbackAsync()).thenReturn(ApiFutures.immediateFuture(null));
- manager.closeAsync();
- }
- }
 }

google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java

@@ -253,6 +253,12 @@ public void runAsync() throws Exception {
  },
  executor);
  assertThat(insertCount.get()).isEqualTo(1L);
+ if (env.getTestHelper().getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()) {
+ // The runAsync() method should only be called once on the runner.
+ // However, due to a bug in regular sessions, it can be executed multiple times on the same
+ // runner.
+ runner = client.runAsync();
+ }
  ApiFuture<Long> deleteCount =
  runner.runAsync(
  txn ->
@@ -299,6 +305,12 @@ public void runAsyncBatchUpdate() throws Exception {
  },
  executor);
  assertThat(insertCount.get()).asList().containsExactly(1L, 1L, 1L);
+ if (env.getTestHelper().getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()) {
+ // The runAsync() method should only be called once on the runner.
+ // However, due to a bug in regular sessions, it can be executed multiple times on the same
+ // runner.
+ runner = client.runAsync();
+ }
  ApiFuture<long[]> deleteCount =
  runner.runAsync(
  txn ->

google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITBatchDmlTest.java

@@ -17,6 +17,7 @@
 package com.google.cloud.spanner.it;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assume.assumeFalse;
 
 import com.google.api.gax.longrunning.OperationFuture;
 import com.google.cloud.spanner.Database;
@@ -84,6 +85,10 @@ public void dropTable() throws Exception {
 
  @Test
  public void noStatementsInRequest() {
+ // TODO(sriharshach): Remove this skip once backend support empty transactions to commit.
+ assumeFalse(
+ "Skipping for multiplexed sessions since it does not allow empty transactions to commit",
+ isUsingMultiplexedSessionsForRW());
  final TransactionCallable<long[]> callable =
  transaction -> {
  List<Statement> stmts = new ArrayList<>();
@@ -252,4 +257,8 @@ public void largeBatchDml_withNonParameterisedStatements() {
  assertThat(actualRowCounts.length).isEqualTo(80);
  assertThat(expectedRowCounts).isEqualTo(actualRowCounts);
  }
+
+ boolean isUsingMultiplexedSessionsForRW() {
+ return env.getTestHelper().getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW();
+ }
 }

google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITBatchReadTest.java

@@ -181,7 +181,9 @@ public static void setUpDatabase() throws Exception {
  totalSize = 0;
  }
  }
- dbClient.write(mutations);
+ if (!mutations.isEmpty()) {
+ dbClient.write(mutations);
+ }
  }
  // Our read/queries are executed with some staleness.
  Thread.sleep(2 * STALENESS_MILLISEC);

google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITTransactionTest.java

@@ -464,7 +464,10 @@ public void nestedSingleUseReadTxnThrows() {
  @Test
  public void nestedTxnSucceedsWhenAllowed() {
  assumeFalse("Emulator does not support multiple parallel transactions", isUsingEmulator());
-
+ // TODO(sriharshach): Remove this skip once backend support empty transactions to commit.
+ assumeFalse(
+ "Skipping for multiplexed sessions since it does not allow empty transactions to commit",
+ isUsingMultiplexedSessionsForRW());
  client
  .readWriteTransaction()
  .allowNestedTransaction()
@@ -588,4 +591,8 @@ public void testTransactionRunnerReturnsCommitStats() {
  // MutationCount = 2 (2 columns).
  assertEquals(2L, runner.getCommitResponse().getCommitStats().getMutationCount());
  }
+
+ boolean isUsingMultiplexedSessionsForRW() {
+ return env.getTestHelper().getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW();
+ }
 }