Skip to content

Commit 5a2f569

Browse files
authored
---
yaml --- r: 33759 b: refs/heads/autosynth-redis c: f2c7d7c h: refs/heads/master i: 33757: c14d11d 33755: aaf0205 33751: 757b3a5 33743: dc1397f 33727: 0c78e0d
1 parent 40fc4b2 commit 5a2f569

3 files changed

Lines changed: 114 additions & 4 deletions

File tree

[refs]

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,7 @@ refs/heads/autosynth-iot: 044be280805a59e06d09658688c9ee474a9815ad
135135
refs/heads/autosynth-kms: d31449d6621a50fb16a4bef4f30f0f3051d27d7c
136136
refs/heads/autosynth-language: 6130869312f99a1e7d3aa0485759172a23333cc5
137137
refs/heads/autosynth-os-login: 49028d40ac477fca5f948cc5a3ce7422729fdb67
138-
refs/heads/autosynth-redis: 1ca609943f2e211420301d543cfd2e8ae059bcc1
138+
refs/heads/autosynth-redis: f2c7d7c0d07091b7729d3e891dfb4354184dc13c
139139
refs/heads/autosynth-scheduler: 57f9fdb1e7de30c85f4ec7198931a07f50603e55
140140
refs/heads/autosynth-spanner: de02ca32edea133b68b51052e325359a3704b5d2
141141
refs/heads/autosynth-speech: 64692f6db11364f663921be02c08072b966b6e7b

branches/autosynth-redis/google-cloud-clients/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,11 @@ class DatabaseClientImpl implements DatabaseClient {
3232
private static final String PARTITION_DML_TRANSACTION = "CloudSpanner.PartitionDMLTransaction";
3333
private static final Tracer tracer = Tracing.getTracer();
3434

35+
private enum SessionMode {
36+
READ,
37+
READ_WRITE
38+
}
39+
3540
static {
3641
TraceUtil.exportSpans(READ_WRITE_TRANSACTION, READ_ONLY_TRANSACTION, PARTITION_DML_TRANSACTION);
3742
}
@@ -57,6 +62,7 @@ public Timestamp write(final Iterable<Mutation> mutations) throws SpannerExcepti
5762
Span span = tracer.spanBuilder(READ_WRITE_TRANSACTION).startSpan();
5863
try (Scope s = tracer.withSpan(span)) {
5964
return runWithSessionRetry(
65+
SessionMode.READ_WRITE,
6066
new Function<Session, Timestamp>() {
6167
@Override
6268
public Timestamp apply(Session session) {
@@ -76,6 +82,7 @@ public Timestamp writeAtLeastOnce(final Iterable<Mutation> mutations) throws Spa
7682
Span span = tracer.spanBuilder(READ_WRITE_TRANSACTION).startSpan();
7783
try (Scope s = tracer.withSpan(span)) {
7884
return runWithSessionRetry(
85+
SessionMode.READ_WRITE,
7986
new Function<Session, Timestamp>() {
8087
@Override
8188
public Timestamp apply(Session session) {
@@ -182,7 +189,10 @@ public TransactionManager transactionManager() {
182189
public long executePartitionedUpdate(final Statement stmt) {
183190
Span span = tracer.spanBuilder(PARTITION_DML_TRANSACTION).startSpan();
184191
try (Scope s = tracer.withSpan(span)) {
192+
// A partitioned update transaction does not need a prepared write session, as the transaction
193+
// object will start a new transaction with specific options anyway.
185194
return runWithSessionRetry(
195+
SessionMode.READ,
186196
new Function<Session, Long>() {
187197
@Override
188198
public Long apply(Session session) {
@@ -195,13 +205,17 @@ public Long apply(Session session) {
195205
}
196206
}
197207

198-
private <T> T runWithSessionRetry(Function<Session, T> callable) {
199-
PooledSession session = getReadWriteSession();
208+
private <T> T runWithSessionRetry(SessionMode mode, Function<Session, T> callable) {
209+
PooledSession session =
210+
mode == SessionMode.READ_WRITE ? getReadWriteSession() : getReadSession();
200211
while (true) {
201212
try {
202213
return callable.apply(session);
203214
} catch (SessionNotFoundException e) {
204-
session = pool.replaceReadWriteSession(e, session);
215+
session =
216+
mode == SessionMode.READ_WRITE
217+
? pool.replaceReadWriteSession(e, session)
218+
: pool.replaceReadSession(e, session);
205219
}
206220
}
207221
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
/*
2+
* Copyright 2019 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.spanner;
18+
19+
import static org.hamcrest.CoreMatchers.equalTo;
20+
import static org.hamcrest.CoreMatchers.is;
21+
import static org.junit.Assert.assertThat;
22+
23+
import com.google.api.gax.grpc.testing.LocalChannelProvider;
24+
import com.google.cloud.NoCredentials;
25+
import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
26+
import io.grpc.Server;
27+
import io.grpc.inprocess.InProcessServerBuilder;
28+
import java.io.IOException;
29+
import org.junit.After;
30+
import org.junit.AfterClass;
31+
import org.junit.Before;
32+
import org.junit.BeforeClass;
33+
import org.junit.Test;
34+
import org.junit.runner.RunWith;
35+
import org.junit.runners.JUnit4;
36+
37+
@RunWith(JUnit4.class)
38+
public class DatabaseClientImplTest {
39+
private static MockSpannerServiceImpl mockSpanner;
40+
private static Server server;
41+
private static LocalChannelProvider channelProvider;
42+
private static Spanner spanner;
43+
private static final Statement UPDATE_STATEMENT =
44+
Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=2");
45+
private static final long UPDATE_COUNT = 1L;
46+
47+
@BeforeClass
48+
public static void startStaticServer() throws IOException {
49+
mockSpanner = new MockSpannerServiceImpl();
50+
mockSpanner.setAbortProbability(0.0D); // We don't want any unpredictable aborted transactions.
51+
mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT));
52+
53+
String uniqueName = InProcessServerBuilder.generateName();
54+
server =
55+
InProcessServerBuilder.forName(uniqueName)
56+
.directExecutor()
57+
.addService(mockSpanner)
58+
.build()
59+
.start();
60+
channelProvider = LocalChannelProvider.create(uniqueName);
61+
}
62+
63+
@AfterClass
64+
public static void stopServer() {
65+
server.shutdown();
66+
}
67+
68+
@Before
69+
public void setUp() throws IOException {
70+
mockSpanner.reset();
71+
spanner =
72+
SpannerOptions.newBuilder()
73+
.setProjectId("[PROJECT]")
74+
.setChannelProvider(channelProvider)
75+
.setCredentials(NoCredentials.getInstance())
76+
.build()
77+
.getService();
78+
}
79+
80+
@After
81+
public void tearDown() throws Exception {
82+
spanner.close();
83+
}
84+
85+
/**
86+
* Test that the update statement can be executed as a partitioned transaction that returns a
87+
* lower bound update count.
88+
*/
89+
@Test
90+
public void testExecutePartitionedDml() {
91+
DatabaseClient client =
92+
spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
93+
long updateCount = client.executePartitionedUpdate(UPDATE_STATEMENT);
94+
assertThat(updateCount, is(equalTo(UPDATE_COUNT)));
95+
}
96+
}

0 commit comments

Comments
 (0)