feat(spanner): Updated unit tests of Partitioned operations for Multiplexed Session.

Author: pratickchokhani

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

@@ -36,7 +36,7 @@ class DatabaseClientImpl implements DatabaseClient {
   @VisibleForTesting final String clientId;
   @VisibleForTesting final SessionPool pool;
   @VisibleForTesting final MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient;
-    final boolean useMultiplexedSessionPartitionedOps;
+  @VisibleForTesting final boolean useMultiplexedSessionPartitionedOps;
   @VisibleForTesting final boolean useMultiplexedSessionForRW;
 
   final boolean useMultiplexedSessionBlindWrite;
@@ -48,7 +48,7 @@ class DatabaseClientImpl implements DatabaseClient {
         pool,
         /* useMultiplexedSessionBlindWrite = */ false,
         /* multiplexedSessionDatabaseClient = */ null,
-            /* useMultiplexedSessionPartitionedOps= */ false,
+        /* useMultiplexedSessionPartitionedOps= */ false,
         tracer,
         /* useMultiplexedSessionForRW = */ false);
   }
@@ -60,7 +60,7 @@ class DatabaseClientImpl implements DatabaseClient {
         pool,
         /* useMultiplexedSessionBlindWrite = */ false,
         /* multiplexedSessionDatabaseClient = */ null,
-            /* useMultiplexedSessionPartitionedOps= */ false,
+        /* useMultiplexedSessionPartitionedOps= */ false,
         tracer,
         /* useMultiplexedSessionForRW = */ false);
   }

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

@@ -23,8 +23,8 @@
 import com.google.cloud.Timestamp;
 import com.google.cloud.spanner.DelayedReadContext.DelayedReadOnlyTransaction;
 import com.google.cloud.spanner.MultiplexedSessionDatabaseClient.MultiplexedSessionTransaction;
-import com.google.cloud.spanner.Options.UpdateOption;
 import com.google.cloud.spanner.Options.TransactionOption;
+import com.google.cloud.spanner.Options.UpdateOption;
 import com.google.common.util.concurrent.MoreExecutors;
 import java.util.concurrent.ExecutionException;
 

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

@@ -360,7 +360,9 @@ private static Boolean getUseMultiplexedSessionFromEnvVariable() {
     return parseBooleanEnvVariable("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS");
   }
 
-  private static Boolean getUseMultiplexedSessionFromEnvVariablePartitionedOps() {
+  @VisibleForTesting
+  @InternalApi
+  protected static Boolean getUseMultiplexedSessionFromEnvVariablePartitionedOps() {
     return parseBooleanEnvVariable("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS");
   }
 

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

@@ -308,7 +308,7 @@ public DatabaseClient getDatabaseClient(DatabaseId db) {
                 pool,
                 getOptions().getSessionPoolOptions().getUseMultiplexedSessionBlindWrite(),
                 multiplexedSessionDatabaseClient,
-                    getOptions().getSessionPoolOptions().getUseMultiplexedSessionPartitionedOps(),
+                getOptions().getSessionPoolOptions().getUseMultiplexedSessionPartitionedOps(),
                 useMultiplexedSessionForRW);
         dbClients.put(db, dbClient);
         return dbClient;
@@ -329,7 +329,7 @@ DatabaseClientImpl createDatabaseClient(
         pool,
         useMultiplexedSessionBlindWrite,
         multiplexedSessionClient,
-            useMultiplexedSessionPartitionedOps,
+        useMultiplexedSessionPartitionedOps,
         tracer,
         useMultiplexedSessionForRW);
   }

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

@@ -56,6 +56,7 @@ public final class BatchClientImplTest {
   private static final String SESSION_NAME = DB_NAME + "/sessions/s1";
   private static final ByteString TXN_ID = ByteString.copyFromUtf8("my-txn");
   private static final String TIMESTAMP = "2017-11-15T10:54:20Z";
+  private static boolean isMultiplexedSession = false;
 
   @Mock private SpannerRpc gapicRpc;
   @Mock private SpannerOptions spannerOptions;
@@ -68,6 +69,11 @@ public final class BatchClientImplTest {
   public static void setupOpenTelemetry() {
     SpannerOptions.resetActiveTracingFramework();
     SpannerOptions.enableOpenTelemetryTraces();
+    Boolean useMultiplexedSessionFromEnvVariablePartitionedOps =
+        SessionPoolOptions.getUseMultiplexedSessionFromEnvVariablePartitionedOps();
+    isMultiplexedSession =
+        useMultiplexedSessionFromEnvVariablePartitionedOps != null
+            && useMultiplexedSessionFromEnvVariablePartitionedOps;
   }
 
   @SuppressWarnings("unchecked")
@@ -88,18 +94,31 @@ public void setUp() {
     when(spannerOptions.getTransportOptions()).thenReturn(transportOptions);
     SessionPoolOptions sessionPoolOptions = mock(SessionPoolOptions.class);
     when(sessionPoolOptions.getPoolMaintainerClock()).thenReturn(Clock.INSTANCE);
+    when(sessionPoolOptions.getUseMultiplexedSessionPartitionedOps())
+        .thenReturn(isMultiplexedSession);
     when(spannerOptions.getSessionPoolOptions()).thenReturn(sessionPoolOptions);
     @SuppressWarnings("resource")
     SpannerImpl spanner = new SpannerImpl(gapicRpc, spannerOptions);
-    client = new BatchClientImpl(spanner.getSessionClient(db), false);
+    client = new BatchClientImpl(spanner.getSessionClient(db), isMultiplexedSession);
   }
 
   @SuppressWarnings("unchecked")
   @Test
   public void testBatchReadOnlyTxnWithBound() throws Exception {
-    Session sessionProto = Session.newBuilder().setName(SESSION_NAME).build();
-    when(gapicRpc.createSession(eq(DB_NAME), anyString(), anyMap(), optionsCaptor.capture()))
-        .thenReturn(sessionProto);
+    Session sessionProto =
+        Session.newBuilder().setName(SESSION_NAME).setMultiplexed(isMultiplexedSession).build();
+    if (isMultiplexedSession) {
+      when(gapicRpc.createSession(
+              eq(DB_NAME),
+              anyString(),
+              anyMap(),
+              optionsCaptor.capture(),
+              eq(isMultiplexedSession)))
+          .thenReturn(sessionProto);
+    } else {
+      when(gapicRpc.createSession(eq(DB_NAME), anyString(), anyMap(), optionsCaptor.capture()))
+          .thenReturn(sessionProto);
+    }
     com.google.protobuf.Timestamp timestamp = Timestamps.parse(TIMESTAMP);
     Transaction txnMetadata =
         Transaction.newBuilder().setId(TXN_ID).setReadTimestamp(timestamp).build();
@@ -113,12 +132,14 @@ public void testBatchReadOnlyTxnWithBound() throws Exception {
     assertThat(batchTxn.getReadTimestamp()).isEqualTo(t);
     assertThat(batchTxn.getReadTimestamp())
         .isEqualTo(batchTxn.getBatchTransactionId().getTimestamp());
+    assertEquals(batchTxn.getBatchTransactionId().isMultiplexedSession(), isMultiplexedSession);
   }
 
   @Test
   public void testBatchReadOnlyTxnWithTxnId() {
     when(txnID.getSessionId()).thenReturn(SESSION_NAME);
     when(txnID.getTransactionId()).thenReturn(TXN_ID);
+    when(txnID.isMultiplexedSession()).thenReturn(isMultiplexedSession);
     Timestamp t = Timestamp.parseTimestamp(TIMESTAMP);
     when(txnID.getTimestamp()).thenReturn(t);
 
@@ -128,6 +149,8 @@ public void testBatchReadOnlyTxnWithTxnId() {
     assertThat(batchTxn.getReadTimestamp()).isEqualTo(t);
     assertThat(batchTxn.getReadTimestamp())
         .isEqualTo(batchTxn.getBatchTransactionId().getTimestamp());
+    assertThat(batchTxn.getBatchTransactionId().isMultiplexedSession())
+        .isEqualTo(isMultiplexedSession);
   }
 
   @Test

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

@@ -95,6 +95,7 @@ public void createSpannerInstance() {
                     .setUseMultiplexedSession(true)
                     .setUseMultiplexedSessionBlindWrite(true)
                     .setUseMultiplexedSessionForRW(true)
+                    .setUseMultiplexedSessionPartitionedOps(true)
                     // Set the maintainer to loop once every 1ms
                     .setMultiplexedSessionMaintenanceLoopFrequency(Duration.ofMillis(1L))
                     // Set multiplexed sessions to be replaced once every 1ms

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

@@ -1273,6 +1273,9 @@ public void transactionManagerReadRowUsingIndexInvalidatedDuringTransaction()
 
   @Test
   public void partitionedDml() throws InterruptedException {
+    assumeFalse(
+        "Multiplexed session do not throw a SessionNotFound errors. ",
+        spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionPartitionedOps());
     assertThrowsSessionNotFoundIfShouldFail(
         () -> client.executePartitionedUpdate(UPDATE_STATEMENT));
   }