Skip to content

feat: add option for multiplexed sessions with partitioned operations #3635

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
Feb 20, 2025
Merged

feat: add option for multiplexed sessions with partitioned operations #3635

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
d56c573
chore(spanner): Release Multiplexed session for Partitioned Ops.
pratickchokhani Feb 11, 2025
e9410cc
Merge branch 'main' into main
pratickchokhani Feb 11, 2025
90d29b1
Merge branch 'main' into main
pratickchokhani Feb 18, 2025
606e272
Merge branch 'main' into main
pratickchokhani Feb 19, 2025
a03d3ea
fix(spanner): Fix unit tests for Partitioned Query and Partitioned DM…
pratickchokhani Feb 20, 2025
c0fe78e
fix(spanner): Fix multiplexed session environment variable validation.
pratickchokhani Feb 20, 2025
235e35d
fix(spanner): Enabled multiplexed session for partitioned ops in inte…
pratickchokhani Feb 20, 2025
3c8ce93
Merge branch 'main' into main
harshachinta Feb 20, 2025
58c9a40
fix(spanner): Disable multiplexed session for partitioned ops in emul…
pratickchokhani Feb 20, 2025
ca1facd
chore(spanner): fix test
harshachinta Feb 20, 2025
8db9d83
chore(spanner): fix test
harshachinta Feb 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 1 addition & 3 deletions google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
1 change: 1 addition & 0 deletions google-cloud-spanner/src/test/java/com/google/cloud/spanner/BatchClientImplTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
8 changes: 8 additions & 0 deletions ...oud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
39 changes: 16 additions & 23 deletions ...ner/src/test/java/com/google/cloud/spanner/connection/PartitionedQueryMockServerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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());
Expand All @@ -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(
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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));
Expand Down Expand Up @@ -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));
}
Expand Down Expand Up @@ -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));
Expand Down
Loading