chore(spanner): verify error message of the exeception

Author: harshachinta

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

@@ -98,7 +98,8 @@ void onError(SpannerException spannerException) {
         this.client.resourceNotFoundException.set((ResourceNotFoundException) spannerException);
       }
       // Mark multiplexed sessions for RW as unimplemented and fall back to regular sessions if
-      // UNIMPLEMENTED is returned.
+      // UNIMPLEMENTED with error message "Transaction type read_write not supported with
+      // multiplexed sessions" is returned.
       this.client.maybeMarkUnimplementedForRW(spannerException);
     }
 
@@ -172,7 +173,6 @@ public void close() {
   /** The current multiplexed session that is used by this client. */
   private final AtomicReference<ApiFuture<SessionReference>> multiplexedSessionReference;
 
-  @VisibleForTesting
   /**
    * The Transaction response returned by the BeginTransaction request with read-write when a
    * multiplexed session is created during client initialization.
@@ -256,7 +256,7 @@ public void onSessionReady(SessionImpl session) {
                 .getOptions()
                 .getSessionPoolOptions()
                 .getUseMultiplexedSessionForRW()) {
-              verifyBeginTransactionWithRWOnMultiplexedSession(session.getName());
+              verifyBeginTransactionWithRWOnMultiplexedSessionAsync(session.getName());
             }
           }
 
@@ -301,14 +301,25 @@ private void maybeMarkUnimplemented(Throwable t) {
 
   private void maybeMarkUnimplementedForRW(SpannerException spannerException) {
     if (spannerException.getErrorCode() == ErrorCode.UNIMPLEMENTED
-    // && spannerException.getCause().getMessage().contains("Transaction type read_write not
-    // supported with multiplexed sessions")
-    ) {
+        && verifyErrorMessage(
+            spannerException,
+            "Transaction type read_write not supported with multiplexed sessions")) {
       unimplementedForRW.set(true);
     }
   }
 
-  private void verifyBeginTransactionWithRWOnMultiplexedSession(String sessionName) {
+  private boolean verifyErrorMessage(SpannerException spannerException, String message) {
+    if (spannerException.getCause() == null) {
+      return false;
+    }
+    if (spannerException.getCause().getMessage() == null) {
+      return false;
+    }
+    return spannerException.getCause().getMessage().contains(message);
+  }
+
+  private void verifyBeginTransactionWithRWOnMultiplexedSessionAsync(String sessionName) {
+    // TODO: Remove once this is guaranteed to be available.
     // annotate the explict BeginTransactionRequest with a transaction tag
     // "multiplexed-rw-background-begin-txn" to avoid storing this request on mock spanner.
     // this is to safeguard other mock spanner tests whose BeginTransaction request count will

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

@@ -1872,6 +1872,7 @@ private Transaction getTemporaryTransactionOrNull(TransactionSelector tx) {
   @Override
   public void beginTransaction(
       BeginTransactionRequest request, StreamObserver<Transaction> responseObserver) {
+    // TODO: Remove once this is guaranteed to be available.
     // Skip storing the explicit BeginTransactionRequest used to verify read-write transaction
     // server availability on multiplexed sessions.
     // This code will be removed once read-write multiplexed sessions are stable on the backend,

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

@@ -1516,6 +1516,65 @@ public void testReadWriteUnimplemented_firstReceivesError_secondFallsBackToRegul
     assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get());
   }
 
+  @Test
+  public void testOtherUnimplementedError_ReadWriteTransactionStillUsesMultiplexedSession() {
+    mockSpanner.setExecuteStreamingSqlExecutionTime(
+        SimulatedExecutionTime.ofException(
+            Status.UNIMPLEMENTED
+                .withDescription("Multiplexed sessions are not supported.")
+                .asRuntimeException()));
+
+    DatabaseClientImpl client =
+        (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+
+    // Wait until the initial BeginTransaction RPC with read-write is complete.
+    assertNotNull(client.multiplexedSessionDatabaseClient);
+    Transaction txn =
+        client.multiplexedSessionDatabaseClient.getReadWriteBeginTransactionReference();
+    assertNotNull(txn);
+    assertNotNull(txn.getId());
+    assertFalse(client.multiplexedSessionDatabaseClient.unimplementedForRW.get());
+
+    // Try to execute a query using single use transaction.
+    try (ResultSet resultSet = client.singleUse().executeQuery(STATEMENT)) {
+      SpannerException spannerException = assertThrows(SpannerException.class, resultSet::next);
+      assertEquals(ErrorCode.UNIMPLEMENTED, spannerException.getErrorCode());
+    }
+    // Verify other UNIMPLEMENTED errors does not turn off read-write transactions to use
+    // multiplexed sessions.
+    assertFalse(client.multiplexedSessionDatabaseClient.unimplementedForRW.get());
+
+    // The read-write transaction should use multiplexed sessions and succeed.
+    client
+        .readWriteTransaction()
+        .run(
+            transaction -> {
+              // Returns a ResultSet containing the precommit token (ResultSetPrecommitToken)
+              transaction.executeUpdate(UPDATE_STATEMENT);
+
+              // Verify that a precommit token is received. This guarantees that the read-write
+              // transaction was executed on a multiplexed session.
+              TransactionContextImpl impl = (TransactionContextImpl) transaction;
+              assertNotNull(impl.getLatestPrecommitToken());
+              assertEquals(
+                  ByteString.copyFromUtf8("ResultSetPrecommitToken"),
+                  impl.getLatestPrecommitToken().getPrecommitToken());
+              return null;
+            });
+
+    // Verify that two ExecuteSqlRequests were received and second one uses a multiplexed session.
+    assertEquals(2, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+
+    Session session2 = mockSpanner.getSession(requests.get(1).getSession());
+    assertNotNull(session2);
+    assertTrue(session2.getMultiplexed());
+
+    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 =