Skip to content

Commit 5afa652

Browse files
harshachintaharshachinta
committed
chore(spanner): fix tests in connect api for mux rw
1 parent 7af512b commit 5afa652

File tree

4 files changed

+100
-20
lines changed

4 files changed

+100
-20
lines changed

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbortedTest.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,8 @@ public void testCommitAborted() {
7272
AbortInterceptor interceptor = new AbortInterceptor(0);
7373
try (ITConnection connection =
7474
createConnection(interceptor, new CountTransactionRetryListener())) {
75+
interceptor.setUsingMultiplexedSession(
76+
isMultiplexedSessionsEnabledForRW(connection.getSpanner()));
7577
// verify that the there is no test record
7678
try (ResultSet rs =
7779
connection.executeQuery(Statement.of("SELECT COUNT(*) AS C FROM TEST WHERE ID=1"))) {
@@ -112,6 +114,8 @@ public void testCommitAbortedDuringUpdateWithReturning() {
112114
AbortInterceptor interceptor = new AbortInterceptor(0);
113115
try (ITConnection connection =
114116
createConnection(interceptor, new CountTransactionRetryListener())) {
117+
interceptor.setUsingMultiplexedSession(
118+
isMultiplexedSessionsEnabledForRW(connection.getSpanner()));
115119
// verify that the there is no test record
116120
try (ResultSet rs =
117121
connection.executeQuery(Statement.of("SELECT COUNT(*) AS C FROM TEST WHERE ID=1"))) {

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -327,4 +327,11 @@ boolean isMultiplexedSessionsEnabled(Spanner spanner) {
327327
}
328328
return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession();
329329
}
330+
331+
boolean isMultiplexedSessionsEnabledForRW(Spanner spanner) {
332+
if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) {
333+
return false;
334+
}
335+
return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW();
336+
}
330337
}

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ITAbstractSpannerTest.java‎

Lines changed: 46 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import com.google.cloud.spanner.GceTestEnvConfig;
2323
import com.google.cloud.spanner.IntegrationTestEnv;
2424
import com.google.cloud.spanner.ResultSet;
25+
import com.google.cloud.spanner.Spanner;
2526
import com.google.cloud.spanner.SpannerExceptionFactory;
2627
import com.google.cloud.spanner.SpannerOptions;
2728
import com.google.cloud.spanner.Statement;
@@ -95,6 +96,8 @@ static ExecutionStep of(StatementExecutionStep step) {
9596
private boolean onlyInjectOnce = false;
9697
private final Random random = new Random();
9798

99+
private boolean usingMultiplexedsession = false;
100+
98101
public AbortInterceptor(double probability) {
99102
Preconditions.checkArgument(probability >= 0.0D && probability <= 1.0D);
100103
this.probability = probability;
@@ -110,6 +113,14 @@ public void setOnlyInjectOnce(boolean value) {
110113
this.onlyInjectOnce = value;
111114
}
112115

116+
/**
117+
* Set this value to true if a multiplexed session is being used. Determining this directly from
118+
* TransactionManagerImpl is challenging as it is a private class.
119+
*/
120+
public void setUsingMultiplexedSession(boolean value) {
121+
this.usingMultiplexedsession = value;
122+
}
123+
113124
protected boolean shouldAbort(String statement, ExecutionStep step) {
114125
return probability > random.nextDouble();
115126
}
@@ -133,27 +144,35 @@ public void intercept(
133144
return;
134145
}
135146
Class<?> cls = Class.forName("com.google.cloud.spanner.TransactionManagerImpl");
136-
Class<?> cls2 =
137-
Class.forName("com.google.cloud.spanner.SessionPool$AutoClosingTransactionManager");
138-
Field delegateField = cls2.getDeclaredField("delegate");
139-
delegateField.setAccessible(true);
140-
watch = watch.reset().start();
141-
while (delegateField.get(tx) == null && watch.elapsed(TimeUnit.MILLISECONDS) < 100) {
142-
Thread.sleep(1L);
147+
if (usingMultiplexedsession) {
148+
Field stateField = cls.getDeclaredField("txnState");
149+
stateField.setAccessible(true);
150+
tx.rollback();
151+
stateField.set(tx, TransactionState.ABORTED);
152+
} else {
153+
Class<?> cls2 =
154+
Class.forName(
155+
"com.google.cloud.spanner.SessionPool$AutoClosingTransactionManager");
156+
Field delegateField = cls2.getDeclaredField("delegate");
157+
delegateField.setAccessible(true);
158+
watch = watch.reset().start();
159+
while (delegateField.get(tx) == null && watch.elapsed(TimeUnit.MILLISECONDS) < 100) {
160+
Thread.sleep(1L);
161+
}
162+
TransactionManager delegate = (TransactionManager) delegateField.get(tx);
163+
if (delegate == null) {
164+
return;
165+
}
166+
Field stateField = cls.getDeclaredField("txnState");
167+
stateField.setAccessible(true);
168+
169+
// First rollback the delegate, and then pretend it aborted.
170+
// We should call rollback on the delegate and not the wrapping
171+
// AutoClosingTransactionManager, as the latter would cause the session to be returned
172+
// to the session pool.
173+
delegate.rollback();
174+
stateField.set(delegate, TransactionState.ABORTED);
143175
}
144-
TransactionManager delegate = (TransactionManager) delegateField.get(tx);
145-
if (delegate == null) {
146-
return;
147-
}
148-
Field stateField = cls.getDeclaredField("txnState");
149-
stateField.setAccessible(true);
150-
151-
// First rollback the delegate, and then pretend it aborted.
152-
// We should call rollback on the delegate and not the wrapping
153-
// AutoClosingTransactionManager, as the latter would cause the session to be returned
154-
// to the session pool.
155-
delegate.rollback();
156-
stateField.set(delegate, TransactionState.ABORTED);
157176
} catch (Exception e) {
158177
throw new RuntimeException(e);
159178
}
@@ -350,4 +369,11 @@ protected boolean indexExists(Connection connection, String table, String index)
350369
}
351370
return false;
352371
}
372+
373+
protected boolean isMultiplexedSessionsEnabledForRW(Spanner spanner) {
374+
if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) {
375+
return false;
376+
}
377+
return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW();
378+
}
353379
}

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/OpenTelemetryTracingTest.java‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import static org.junit.Assert.assertEquals;
2121
import static org.junit.Assert.assertFalse;
2222
import static org.junit.Assert.assertTrue;
23+
import static org.junit.Assume.assumeFalse;
2324

2425
import com.google.cloud.spanner.MockSpannerServiceImpl;
2526
import com.google.cloud.spanner.ResultSet;
@@ -111,6 +112,9 @@ Connection createTestConnection(String url) {
111112
public void testSingleUseQuery_withoutSqlStatement() {
112113
try (Connection connection = createTestConnection(getBaseUrl())) {
113114
connection.setAutocommit(true);
115+
assumeFalse(
116+
"OpenTelemetryTracingTest handler is not implemented for read-write with multiplexed sessions",
117+
isMultiplexedSessionsEnabledForRW(connection.getSpanner()));
114118
try (ResultSet resultSet = connection.executeQuery(SELECT1_STATEMENT)) {
115119
assertTrue(resultSet.next());
116120
assertFalse(resultSet.next());
@@ -149,6 +153,9 @@ public boolean isEnableExtendedTracing() {
149153

150154
try (Connection connection = createTestConnection(getBaseUrl())) {
151155
connection.setAutocommit(true);
156+
assumeFalse(
157+
"OpenTelemetryTracingTest handler is not implemented for read-write with multiplexed sessions",
158+
isMultiplexedSessionsEnabledForRW(connection.getSpanner()));
152159
try (ResultSet resultSet = connection.executeQuery(SELECT1_STATEMENT)) {
153160
assertTrue(resultSet.next());
154161
assertFalse(resultSet.next());
@@ -185,6 +192,9 @@ public boolean isEnableExtendedTracing() {
185192
public void testSingleUseQuery() {
186193
try (Connection connection = createTestConnection()) {
187194
connection.setAutocommit(true);
195+
assumeFalse(
196+
"OpenTelemetryTracingTest handler is not implemented for read-write with multiplexed sessions",
197+
isMultiplexedSessionsEnabledForRW(connection.getSpanner()));
188198
try (ResultSet resultSet = connection.executeQuery(SELECT1_STATEMENT)) {
189199
assertTrue(resultSet.next());
190200
assertFalse(resultSet.next());
@@ -220,6 +230,9 @@ public void testSingleUseUpdate() {
220230
try (Connection connection = createTestConnection()) {
221231
connection.setAutocommit(true);
222232
connection.executeUpdate(INSERT_STATEMENT);
233+
assumeFalse(
234+
"OpenTelemetryTracingTest handler is not implemented for read-write with multiplexed sessions",
235+
isMultiplexedSessionsEnabledForRW(connection.getSpanner()));
223236
}
224237
assertEquals(CompletableResultCode.ofSuccess(), spanExporter.flush());
225238
List<SpanData> spans = spanExporter.getFinishedSpanItems();
@@ -256,6 +269,9 @@ public void testSingleUseBatchUpdate() {
256269
connection.executeUpdate(INSERT_STATEMENT);
257270
connection.executeUpdate(INSERT_STATEMENT);
258271
connection.runBatch();
272+
assumeFalse(
273+
"OpenTelemetryTracingTest handler is not implemented for read-write with multiplexed sessions",
274+
isMultiplexedSessionsEnabledForRW(connection.getSpanner()));
259275
}
260276
assertEquals(CompletableResultCode.ofSuccess(), spanExporter.flush());
261277
List<SpanData> spans = spanExporter.getFinishedSpanItems();
@@ -297,6 +313,9 @@ public void testSingleUseDdl() {
297313

298314
try (Connection connection = createTestConnection()) {
299315
connection.setAutocommit(true);
316+
assumeFalse(
317+
"OpenTelemetryTracingTest handler is not implemented for read-write with multiplexed sessions",
318+
isMultiplexedSessionsEnabledForRW(connection.getSpanner()));
300319
connection.execute(Statement.of(ddl));
301320
}
302321
assertEquals(CompletableResultCode.ofSuccess(), spanExporter.flush());
@@ -315,6 +334,9 @@ public void testSingleUseDdlBatch() {
315334

316335
try (Connection connection = createTestConnection()) {
317336
connection.setAutocommit(true);
337+
assumeFalse(
338+
"OpenTelemetryTracingTest handler is not implemented for read-write with multiplexed sessions",
339+
isMultiplexedSessionsEnabledForRW(connection.getSpanner()));
318340
connection.startBatchDdl();
319341
connection.execute(Statement.of(ddl1));
320342
connection.execute(Statement.of(ddl2));
@@ -332,6 +354,9 @@ public void testSingleUseDdlBatch() {
332354
public void testMultiUseReadOnlyQueries() {
333355
try (Connection connection = createTestConnection()) {
334356
connection.setAutocommit(false);
357+
assumeFalse(
358+
"OpenTelemetryTracingTest handler is not implemented for read-write with multiplexed sessions",
359+
isMultiplexedSessionsEnabledForRW(connection.getSpanner()));
335360
connection.setReadOnly(true);
336361
twice(
337362
() -> {
@@ -363,6 +388,9 @@ public void testMultiUseReadOnlyQueries() {
363388
public void testMultiUseReadWriteQueries() {
364389
try (Connection connection = createTestConnection()) {
365390
connection.setAutocommit(false);
391+
assumeFalse(
392+
"OpenTelemetryTracingTest handler is not implemented for read-write with multiplexed sessions",
393+
isMultiplexedSessionsEnabledForRW(connection.getSpanner()));
366394
connection.setReadOnly(false);
367395
twice(
368396
() -> {
@@ -397,6 +425,9 @@ public void testMultiUseReadWriteQueries() {
397425
public void testMultiUseReadWriteUpdates() {
398426
try (Connection connection = createTestConnection()) {
399427
connection.setAutocommit(false);
428+
assumeFalse(
429+
"OpenTelemetryTracingTest handler is not implemented for read-write with multiplexed sessions",
430+
isMultiplexedSessionsEnabledForRW(connection.getSpanner()));
400431
connection.setReadOnly(false);
401432
assertEquals(1L, connection.executeUpdate(INSERT_STATEMENT));
402433
assertEquals(1L, connection.executeUpdate(INSERT_STATEMENT));
@@ -426,6 +457,9 @@ public void testMultiUseReadWriteUpdates() {
426457
public void testMultiUseReadWriteBatchUpdates() {
427458
try (Connection connection = createTestConnection()) {
428459
connection.setAutocommit(false);
460+
assumeFalse(
461+
"OpenTelemetryTracingTest handler is not implemented for read-write with multiplexed sessions",
462+
isMultiplexedSessionsEnabledForRW(connection.getSpanner()));
429463
connection.setReadOnly(false);
430464

431465
twice(
@@ -466,6 +500,9 @@ public void testMultiUseReadWriteBatchUpdates() {
466500
public void testMultiUseReadWriteAborted() {
467501
try (Connection connection = createTestConnection()) {
468502
connection.setAutocommit(false);
503+
assumeFalse(
504+
"OpenTelemetryTracingTest handler is not implemented for read-write with multiplexed sessions",
505+
isMultiplexedSessionsEnabledForRW(connection.getSpanner()));
469506
connection.setReadOnly(false);
470507
assertEquals(1L, connection.executeUpdate(INSERT_STATEMENT));
471508
mockSpanner.abortNextStatement();
@@ -514,6 +551,9 @@ public void testSavepoint() {
514551

515552
try (Connection connection = createTestConnection()) {
516553
connection.setAutocommit(false);
554+
assumeFalse(
555+
"OpenTelemetryTracingTest handler is not implemented for read-write with multiplexed sessions",
556+
isMultiplexedSessionsEnabledForRW(connection.getSpanner()));
517557
connection.setReadOnly(false);
518558
connection.setSavepointSupport(SavepointSupport.ENABLED);
519559
assertEquals(1L, connection.executeUpdate(statement1));
@@ -563,6 +603,9 @@ public void testTransactionTag() {
563603
try (Connection connection = createTestConnection()) {
564604
connection.setAutocommit(false);
565605
connection.setReadOnly(false);
606+
assumeFalse(
607+
"OpenTelemetryTracingTest handler is not implemented for read-write with multiplexed sessions",
608+
isMultiplexedSessionsEnabledForRW(connection.getSpanner()));
566609
connection.setTransactionTag("my_tag");
567610
assertEquals(1L, connection.executeUpdate(INSERT_STATEMENT));
568611
connection.commit();

0 commit comments

Comments
 (0)