From d56c573c4c1b2c9075779b1d5f56f3d18de3bbcf Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Tue, 11 Feb 2025 17:09:37 +0530 Subject: [PATCH 1/7] chore(spanner): Release Multiplexed session for Partitioned Ops. --- .../java/com/google/cloud/spanner/SessionPoolOptions.java | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java index 03551640b43..58c0eca8100 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java @@ -370,9 +370,7 @@ private static Boolean getUseMultiplexedSessionFromEnvVariable() { @VisibleForTesting @InternalApi protected static Boolean getUseMultiplexedSessionFromEnvVariablePartitionedOps() { - // Checks the value of env, GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS - // This returns null until Partitioned Operations is supported. - return null; + return parseBooleanEnvVariable("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS"); } private static Boolean parseBooleanEnvVariable(String variableName) { From a03d3eafa1deae1a672ec57cfaec8468185d6450 Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Thu, 20 Feb 2025 12:05:15 +0530 Subject: [PATCH 2/7] fix(spanner): Fix unit tests for Partitioned Query and Partitioned DML Multiplexed support launch. --- .../cloud/spanner/BatchClientImplTest.java | 1 + .../connection/AbstractMockServerTest.java | 8 ++++ .../PartitionedQueryMockServerTest.java | 39 ++++++++----------- 3 files changed, 25 insertions(+), 23 deletions(-) diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BatchClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BatchClientImplTest.java index edafc7ddba9..ba508ac8d19 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BatchClientImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BatchClientImplTest.java @@ -102,6 +102,7 @@ public void setUp() { @SuppressWarnings("resource") SpannerImpl spanner = new SpannerImpl(gapicRpc, spannerOptions); client = new BatchClientImpl(spanner.getSessionClient(db), isMultiplexedSession); + BatchClientImpl.unimplementedForPartitionedOps.set(false); } @SuppressWarnings("unchecked") diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java index ac6a6ecbc77..0ab89b7a9ea 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java @@ -328,6 +328,14 @@ boolean isMultiplexedSessionsEnabled(Spanner spanner) { return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession(); } + boolean isMultiplexedSessionsEnabledForPartitionedOps(Spanner spanner) { + if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) { + return false; + } + return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionPartitionedOps() + && spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession(); + } + boolean isMultiplexedSessionsEnabledForRW(Spanner spanner) { if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) { return false; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/PartitionedQueryMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/PartitionedQueryMockServerTest.java index 655ca0de586..9965b14d3db 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/PartitionedQueryMockServerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/PartitionedQueryMockServerTest.java @@ -93,11 +93,7 @@ public void testPartitionQuery() { assertFalse(resultSet.next()); } } - if (isMultiplexedSessionsEnabled(connection.getSpanner())) { - assertEquals(3, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); - } else { - assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); - } + assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); } assertEquals(2, mockSpanner.countRequestsOfType(BeginTransactionRequest.class)); assertEquals(2, mockSpanner.countRequestsOfType(PartitionQueryRequest.class)); @@ -155,11 +151,7 @@ public void testMixNormalAndPartitionQueryInReadOnlyTransaction() { readTimestamps.add(connection.getReadTimestamp()); connection.commit(); } - if (isMultiplexedSessionsEnabled(connection.getSpanner())) { - assertEquals(3, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); - } else { - assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); - } + assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); } // The above will start two transactions: // 1. The initial 'normal' read-only transaction. @@ -228,6 +220,10 @@ public void testRunPartition() { if (createSessionRequestCounts == expectedCreateSessionsRPC + 1) { isMultiplexedSessionCreated = true; } + } else if (isMultiplexedSessionsEnabledForPartitionedOps(connection.getSpanner()) + && isMultiplexedSessionCreated) { + // When multiplexed session will be reused for each iteration. + assertEquals(0, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); } else { assertEquals( expectedCreateSessionsRPC, @@ -261,6 +257,7 @@ public void testRunPartitionUsingSql() { String prefix = dialect == Dialect.POSTGRESQL ? "spanner." : ""; int maxPartitions = 5; + boolean isMultiplexedSessionCreated = false; try (Connection connection = createConnection()) { connection.execute(Statement.of("set autocommit=true")); assertTrue(connection.isAutocommit()); @@ -284,7 +281,6 @@ public void testRunPartitionUsingSql() { assertFalse(resultSet.next()); } - boolean isMultiplexedSessionCreated = false; for (boolean useLiteral : new boolean[] {true, false}) { try (ResultSet partitions = connection.executeQuery( @@ -328,6 +324,10 @@ public void testRunPartitionUsingSql() { if (createSessionRequestCounts == expectedCreateSessionsRPC + 1) { isMultiplexedSessionCreated = true; } + } else if (isMultiplexedSessionsEnabledForPartitionedOps(connection.getSpanner()) + && isMultiplexedSessionCreated) { + // When multiplexed session will be reused for each iteration. + assertEquals(0, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); } else { assertEquals( expectedCreateSessionsRPC, @@ -570,11 +570,7 @@ public void testRunPartitionedQueryUsingSql() { assertEquals(maxPartitions * generatedRowCount, rowCount); } } - if (isMultiplexedSessionsEnabled(connection.getSpanner())) { - assertEquals(3, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); - } else { - assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); - } + assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); } // We have 2 requests of each, as we run the query with data boost both enabled and disabled. assertEquals(2, mockSpanner.countRequestsOfType(BeginTransactionRequest.class)); @@ -679,8 +675,8 @@ public void testRunPartitionedQueryWithMaxParallelism() { assertEquals(maxPartitions * generatedRowCount, rowCount); } } - if (isMultiplexedSessionsEnabled(connection.getSpanner())) { - assertEquals(6, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + if (isMultiplexedSessionsEnabledForPartitionedOps(connection.getSpanner())) { + assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); } else { assertEquals(5, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); } @@ -758,11 +754,8 @@ public void testAutoPartitionMode() { exception .getMessage() .contains("Partition query is not supported for read/write transaction")); - if (isMultiplexedSessionsEnabled(connection.getSpanner())) { - assertEquals(3, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); - } else { - assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); - } + + assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); } assertEquals(2, mockSpanner.countRequestsOfType(BeginTransactionRequest.class)); assertEquals(2, mockSpanner.countRequestsOfType(PartitionQueryRequest.class)); From c0fe78ea350a98b391cf0ef3cd085be4af1a2b53 Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Thu, 20 Feb 2025 13:21:41 +0530 Subject: [PATCH 3/7] fix(spanner): Fix multiplexed session environment variable validation. --- .../main/java/com/google/cloud/spanner/SessionPoolOptions.java | 2 +- .../cloud/spanner/connection/AbstractMockServerTest.java | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java index 58c0eca8100..71e07996519 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java @@ -360,7 +360,7 @@ public boolean getUseMultiplexedSessionForRW() { @VisibleForTesting @InternalApi public boolean getUseMultiplexedSessionPartitionedOps() { - return useMultiplexedSessionForPartitionedOps; + return getUseMultiplexedSession() && useMultiplexedSessionForPartitionedOps; } private static Boolean getUseMultiplexedSessionFromEnvVariable() { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java index 0ab89b7a9ea..a9b27f21545 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java @@ -332,8 +332,7 @@ boolean isMultiplexedSessionsEnabledForPartitionedOps(Spanner spanner) { if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) { return false; } - return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionPartitionedOps() - && spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession(); + return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionPartitionedOps(); } boolean isMultiplexedSessionsEnabledForRW(Spanner spanner) { From 235e35d7bddae2c480d4dccf6f22eceb2e2afbc4 Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Thu, 20 Feb 2025 15:09:48 +0530 Subject: [PATCH 4/7] fix(spanner): Enabled multiplexed session for partitioned ops in integration test. --- ...egration-tests-against-emulator-with-multiplexed-session.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/integration-tests-against-emulator-with-multiplexed-session.yaml b/.github/workflows/integration-tests-against-emulator-with-multiplexed-session.yaml index bd7dfef3972..b45e489599a 100644 --- a/.github/workflows/integration-tests-against-emulator-with-multiplexed-session.yaml +++ b/.github/workflows/integration-tests-against-emulator-with-multiplexed-session.yaml @@ -40,3 +40,4 @@ jobs: JOB_TYPE: test SPANNER_EMULATOR_HOST: localhost:9010 GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS: true + GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS: true From 58c9a406e2fa4a9d0b4883d4faa8f3c85c972d98 Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Thu, 20 Feb 2025 16:15:34 +0530 Subject: [PATCH 5/7] fix(spanner): Disable multiplexed session for partitioned ops in emulator integration test. --- ...ration-tests-against-emulator-with-multiplexed-session.yaml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.github/workflows/integration-tests-against-emulator-with-multiplexed-session.yaml b/.github/workflows/integration-tests-against-emulator-with-multiplexed-session.yaml index b45e489599a..6dda467b64e 100644 --- a/.github/workflows/integration-tests-against-emulator-with-multiplexed-session.yaml +++ b/.github/workflows/integration-tests-against-emulator-with-multiplexed-session.yaml @@ -39,5 +39,4 @@ jobs: env: JOB_TYPE: test SPANNER_EMULATOR_HOST: localhost:9010 - GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS: true - GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS: true + GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS: true \ No newline at end of file From ca1facdb478ad4ebb8023cb951af4f857ab36a78 Mon Sep 17 00:00:00 2001 From: Sri Harsha CH Date: Thu, 20 Feb 2025 17:23:31 +0000 Subject: [PATCH 6/7] chore(spanner): fix test --- .../PartitionedQueryMockServerTest.java | 32 ++++++++++++++++--- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/PartitionedQueryMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/PartitionedQueryMockServerTest.java index 9965b14d3db..c409a88007a 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/PartitionedQueryMockServerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/PartitionedQueryMockServerTest.java @@ -93,7 +93,13 @@ public void testPartitionQuery() { assertFalse(resultSet.next()); } } - assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + if (isMultiplexedSessionsEnabledForPartitionedOps(connection.getSpanner())) { + assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + } else if (isMultiplexedSessionsEnabled(connection.getSpanner())) { + assertEquals(3, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + } else { + assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + } } assertEquals(2, mockSpanner.countRequestsOfType(BeginTransactionRequest.class)); assertEquals(2, mockSpanner.countRequestsOfType(PartitionQueryRequest.class)); @@ -151,7 +157,13 @@ public void testMixNormalAndPartitionQueryInReadOnlyTransaction() { readTimestamps.add(connection.getReadTimestamp()); connection.commit(); } - assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + if (isMultiplexedSessionsEnabledForPartitionedOps(connection.getSpanner())) { + assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + } else if (isMultiplexedSessionsEnabled(connection.getSpanner())) { + assertEquals(3, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + } else { + assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + } } // The above will start two transactions: // 1. The initial 'normal' read-only transaction. @@ -570,7 +582,13 @@ public void testRunPartitionedQueryUsingSql() { assertEquals(maxPartitions * generatedRowCount, rowCount); } } - assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + if (isMultiplexedSessionsEnabledForPartitionedOps(connection.getSpanner())) { + assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + } else if (isMultiplexedSessionsEnabled(connection.getSpanner())) { + assertEquals(3, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + } else { + assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + } } // We have 2 requests of each, as we run the query with data boost both enabled and disabled. assertEquals(2, mockSpanner.countRequestsOfType(BeginTransactionRequest.class)); @@ -755,7 +773,13 @@ public void testAutoPartitionMode() { .getMessage() .contains("Partition query is not supported for read/write transaction")); - assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + if (isMultiplexedSessionsEnabledForPartitionedOps(connection.getSpanner())) { + assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + } else if (isMultiplexedSessionsEnabled(connection.getSpanner())) { + assertEquals(3, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + } else { + assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + } } assertEquals(2, mockSpanner.countRequestsOfType(BeginTransactionRequest.class)); assertEquals(2, mockSpanner.countRequestsOfType(PartitionQueryRequest.class)); From 8db9d831c7f945df65d4e3f8ee110b99c8917c79 Mon Sep 17 00:00:00 2001 From: Sri Harsha CH Date: Thu, 20 Feb 2025 17:30:42 +0000 Subject: [PATCH 7/7] chore(spanner): fix test --- .../spanner/connection/PartitionedQueryMockServerTest.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/PartitionedQueryMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/PartitionedQueryMockServerTest.java index c409a88007a..cdd8b15a38a 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/PartitionedQueryMockServerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/PartitionedQueryMockServerTest.java @@ -695,6 +695,8 @@ public void testRunPartitionedQueryWithMaxParallelism() { } if (isMultiplexedSessionsEnabledForPartitionedOps(connection.getSpanner())) { assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + } else if (isMultiplexedSessionsEnabled(connection.getSpanner())) { + assertEquals(6, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); } else { assertEquals(5, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); }