chore(spanner): add method to get mux for RW

Author: harshachinta

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

@@ -36,26 +36,29 @@ class DatabaseClientImpl implements DatabaseClient {
   @VisibleForTesting final String clientId;
   @VisibleForTesting final SessionPool pool;
   @VisibleForTesting final MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient;
+  @VisibleForTesting final boolean useMultiplexedSessionForRW;
 
   @VisibleForTesting
   DatabaseClientImpl(SessionPool pool, TraceWrapper tracer) {
-    this("", pool, /* multiplexedSessionDatabaseClient = */ null, tracer);
+    this("", pool, /* multiplexedSessionDatabaseClient = */ null, tracer, false);
   }
 
   @VisibleForTesting
   DatabaseClientImpl(String clientId, SessionPool pool, TraceWrapper tracer) {
-    this(clientId, pool, /* multiplexedSessionDatabaseClient = */ null, tracer);
+    this(clientId, pool, /* multiplexedSessionDatabaseClient = */ null, tracer, false);
   }
 
   DatabaseClientImpl(
       String clientId,
       SessionPool pool,
       @Nullable MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient,
-      TraceWrapper tracer) {
+      TraceWrapper tracer,
+      boolean useMultiplexedSessionForRW) {
     this.clientId = clientId;
     this.pool = pool;
     this.multiplexedSessionDatabaseClient = multiplexedSessionDatabaseClient;
     this.tracer = tracer;
+    this.useMultiplexedSessionForRW = useMultiplexedSessionForRW;
   }
 
   @VisibleForTesting
@@ -72,6 +75,14 @@ DatabaseClient getMultiplexedSession() {
     return pool.getMultiplexedSessionWithFallback();
   }
 
+  @VisibleForTesting
+  DatabaseClient getMultiplexedSessionForRW() {
+    if (this.useMultiplexedSessionForRW) {
+      return getMultiplexedSession();
+    }
+    return getSession();
+  }
+
   @Override
   public Dialect getDialect() {
     return pool.getDialect();
@@ -215,7 +226,7 @@ public ReadOnlyTransaction readOnlyTransaction(TimestampBound bound) {
   public TransactionRunner readWriteTransaction(TransactionOption... options) {
     ISpan span = tracer.spanBuilder(READ_WRITE_TRANSACTION, options);
     try (IScope s = tracer.withSpan(span)) {
-      return getSession().readWriteTransaction(options);
+      return getMultiplexedSessionForRW().readWriteTransaction(options);
     } catch (RuntimeException e) {
       span.setStatus(e);
       span.end();

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

@@ -303,7 +303,7 @@ public DatabaseClient getDatabaseClient(DatabaseId db) {
                 numMultiplexedSessionsReleased);
         pool.maybeWaitOnMinSessions();
         DatabaseClientImpl dbClient =
-            createDatabaseClient(clientId, pool, multiplexedSessionDatabaseClient);
+            createDatabaseClient(clientId, pool, multiplexedSessionDatabaseClient, useMultiplexedSessionForRW);
         dbClients.put(db, dbClient);
         return dbClient;
       }
@@ -314,8 +314,9 @@ public DatabaseClient getDatabaseClient(DatabaseId db) {
   DatabaseClientImpl createDatabaseClient(
       String clientId,
       SessionPool pool,
-      @Nullable MultiplexedSessionDatabaseClient multiplexedSessionClient) {
-    return new DatabaseClientImpl(clientId, pool, multiplexedSessionClient, tracer);
+      @Nullable MultiplexedSessionDatabaseClient multiplexedSessionClient,
+      boolean useMultiplexedSessionForRW) {
+    return new DatabaseClientImpl(clientId, pool, multiplexedSessionClient, tracer, useMultiplexedSessionForRW);
   }
 
   @Override

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

@@ -47,7 +47,7 @@ private static class SpannerWithClosedSessionsImpl extends SpannerImpl {
 
     @Override
     DatabaseClientImpl createDatabaseClient(
-        String clientId, SessionPool pool, MultiplexedSessionDatabaseClient ignore) {
+        String clientId, SessionPool pool, MultiplexedSessionDatabaseClient ignore, boolean useMultiplexedSessionForRW) {
       return new DatabaseClientWithClosedSessionImpl(clientId, pool, tracer);
     }
   }