From 1234447f25e62dbd2a002cfed5a2d6e905e0e617 Mon Sep 17 00:00:00 2001
From: Sri Harsha CH <57220027+harshachinta@users.noreply.github.com>
Date: Fri, 6 Dec 2024 10:57:08 +0530
Subject: [PATCH] =?UTF-8?q?chore(spanner):=20support=20multiplexed=20sessi?=
 =?UTF-8?q?on=20for=20rw=20transactions=20in=20ex=E2=80=A6=20(#3471)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* chore(spanner): support multiplexed session for rw transactions in executor

* chore(spanner): lint fix
---
 .../cloud/executor/spanner/CloudClientExecutor.java   | 11 ++++++++---
 .../cloud/spanner/SessionPoolOptionsHelper.java       |  6 ++++++
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/CloudClientExecutor.java b/google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/CloudClientExecutor.java
index 69c0d4d3ea9..c82b6306eb8 100644
--- a/google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/CloudClientExecutor.java
+++ b/google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/CloudClientExecutor.java
@@ -829,10 +829,15 @@ private synchronized Spanner getClient(long timeoutSeconds, boolean useMultiplex
 
     com.google.cloud.spanner.SessionPoolOptions.Builder poolOptionsBuilder =
         com.google.cloud.spanner.SessionPoolOptions.newBuilder();
-    SessionPoolOptionsHelper.setUseMultiplexedSession(
-        com.google.cloud.spanner.SessionPoolOptions.newBuilder(), useMultiplexedSession);
+    SessionPoolOptionsHelper.setUseMultiplexedSession(poolOptionsBuilder, useMultiplexedSession);
     SessionPoolOptionsHelper.setUseMultiplexedSessionBlindWrite(
-        com.google.cloud.spanner.SessionPoolOptions.newBuilder(), useMultiplexedSession);
+        poolOptionsBuilder, useMultiplexedSession);
+    SessionPoolOptionsHelper.setUseMultiplexedSessionForRW(
+        poolOptionsBuilder, useMultiplexedSession);
+    LOGGER.log(
+        Level.INFO,
+        String.format(
+            "Using multiplexed sessions for read-write transactions: %s", useMultiplexedSession));
     com.google.cloud.spanner.SessionPoolOptions sessionPoolOptions = poolOptionsBuilder.build();
     // Cloud Spanner Client does not support global retry settings,
     // Thus, we need to add retry settings to each individual stub.
diff --git a/google-cloud-spanner-executor/src/main/java/com/google/cloud/spanner/SessionPoolOptionsHelper.java b/google-cloud-spanner-executor/src/main/java/com/google/cloud/spanner/SessionPoolOptionsHelper.java
index dafaa4a1f31..8e085947711 100644
--- a/google-cloud-spanner-executor/src/main/java/com/google/cloud/spanner/SessionPoolOptionsHelper.java
+++ b/google-cloud-spanner-executor/src/main/java/com/google/cloud/spanner/SessionPoolOptionsHelper.java
@@ -38,4 +38,10 @@ public static SessionPoolOptions.Builder setUseMultiplexedSessionBlindWrite(
     return sessionPoolOptionsBuilder.setUseMultiplexedSessionBlindWrite(
         useMultiplexedSessionBlindWrite);
   }
+
+  // TODO: Remove when multiplexed session for read write is released.
+  public static SessionPoolOptions.Builder setUseMultiplexedSessionForRW(
+      SessionPoolOptions.Builder sessionPoolOptionsBuilder, boolean useMultiplexedSessionForRW) {
+    return sessionPoolOptionsBuilder.setUseMultiplexedSessionForRW(useMultiplexedSessionForRW);
+  }
 }