Skip to content

Commit d5ed3f2

Browse files
yihanzhenyihanzhen
committed
Make all the unit tests work. ITs are failing pending on another PR.
1 parent 75d1d0b commit d5ed3f2

File tree

6 files changed

+193
-64
lines changed

6 files changed

+193
-64
lines changed

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java‎

Lines changed: 47 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import com.google.api.client.util.BackOff;
2626
import com.google.api.client.util.ExponentialBackOff;
2727
import com.google.api.gax.paging.Page;
28+
import com.google.api.gax.rpc.ServerStream;
2829
import com.google.api.pathtemplate.PathTemplate;
2930
import com.google.cloud.BaseService;
3031
import com.google.cloud.ByteArray;
@@ -828,7 +829,7 @@ public ReadContext singleUse() {
828829

829830
@Override
830831
public ReadContext singleUse(TimestampBound bound) {
831-
return setActive(new SingleReadContext(this, bound, rawGrpcRpc, defaultPrefetchChunks));
832+
return setActive(new SingleReadContext(this, bound, gapicRpc, defaultPrefetchChunks));
832833
}
833834

834835
@Override
@@ -839,7 +840,7 @@ public ReadOnlyTransaction singleUseReadOnlyTransaction() {
839840
@Override
840841
public ReadOnlyTransaction singleUseReadOnlyTransaction(TimestampBound bound) {
841842
return setActive(
842-
new SingleUseReadOnlyTransaction(this, bound, rawGrpcRpc, defaultPrefetchChunks));
843+
new SingleUseReadOnlyTransaction(this, bound, gapicRpc, defaultPrefetchChunks));
843844
}
844845

845846
@Override
@@ -850,12 +851,12 @@ public ReadOnlyTransaction readOnlyTransaction() {
850851
@Override
851852
public ReadOnlyTransaction readOnlyTransaction(TimestampBound bound) {
852853
return setActive(
853-
new MultiUseReadOnlyTransaction(this, bound, rawGrpcRpc, defaultPrefetchChunks));
854+
new MultiUseReadOnlyTransaction(this, bound, gapicRpc, defaultPrefetchChunks));
854855
}
855856

856857
@Override
857858
public TransactionRunner readWriteTransaction() {
858-
return setActive(new TransactionRunnerImpl(this, rawGrpcRpc, defaultPrefetchChunks));
859+
return setActive(new TransactionRunnerImpl(this, gapicRpc, defaultPrefetchChunks));
859860
}
860861

861862
@Override
@@ -1055,20 +1056,14 @@ ResultSet executeQueryInternalWithOptions(
10551056
new ResumableStreamIterator(MAX_BUFFERED_CHUNKS, QUERY) {
10561057
@Override
10571058
CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
1058-
GrpcStreamIterator stream = new GrpcStreamIterator(prefetchChunks);
1059-
SpannerRpc.StreamingCall call =
1060-
rpc.executeQuery(
1061-
resumeToken == null
1062-
? request
1063-
: request.toBuilder().setResumeToken(resumeToken).build(),
1064-
stream.consumer(),
1065-
session.options);
1066-
// We get one message for free.
1067-
if (prefetchChunks > 1) {
1068-
call.request(prefetchChunks - 1);
1069-
}
1070-
stream.setCall(call);
1071-
return stream;
1059+
return new CloseableServerStreamIterator<PartialResultSet>(rpc.executeQuery(
1060+
resumeToken == null
1061+
? request
1062+
: request.toBuilder().setResumeToken(resumeToken).build(),
1063+
null,
1064+
session.options));
1065+
1066+
// let resume fail for now
10721067
}
10731068
};
10741069
return new GrpcResultSet(stream, this, queryMode);
@@ -1168,20 +1163,14 @@ ResultSet readInternalWithOptions(
11681163
new ResumableStreamIterator(MAX_BUFFERED_CHUNKS, READ) {
11691164
@Override
11701165
CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
1171-
GrpcStreamIterator stream = new GrpcStreamIterator(prefetchChunks);
1172-
SpannerRpc.StreamingCall call =
1173-
rpc.read(
1174-
resumeToken == null
1175-
? request
1176-
: request.toBuilder().setResumeToken(resumeToken).build(),
1177-
stream.consumer(),
1178-
session.options);
1179-
// We get one message for free.
1180-
if (prefetchChunks > 1) {
1181-
call.request(prefetchChunks - 1);
1182-
}
1183-
stream.setCall(call);
1184-
return stream;
1166+
return new CloseableServerStreamIterator<PartialResultSet>(rpc.read(
1167+
resumeToken == null
1168+
? request
1169+
: request.toBuilder().setResumeToken(resumeToken).build(),
1170+
null,
1171+
session.options));
1172+
1173+
// let resume fail for now
11851174
}
11861175
};
11871176
GrpcResultSet resultSet =
@@ -2287,6 +2276,32 @@ interface CloseableIterator<T> extends Iterator<T> {
22872276
void close(@Nullable String message);
22882277
}
22892278

2279+
private static final class CloseableServerStreamIterator<T> implements CloseableIterator<T> {
2280+
2281+
private final ServerStream<T> stream;
2282+
private final Iterator<T> iterator;
2283+
2284+
public CloseableServerStreamIterator(ServerStream<T> stream) {
2285+
this.stream = stream;
2286+
this.iterator = stream.iterator();
2287+
}
2288+
2289+
@Override
2290+
public boolean hasNext() {
2291+
return iterator.hasNext();
2292+
}
2293+
2294+
@Override
2295+
public T next() {
2296+
return iterator.next();
2297+
}
2298+
2299+
@Override
2300+
public void close(@Nullable String message) {
2301+
stream.cancel();
2302+
}
2303+
}
2304+
22902305
/** Adapts a streaming read/query call into an iterator over partial result sets. */
22912306
@VisibleForTesting
22922307
static class GrpcStreamIterator extends AbstractIterator<PartialResultSet>

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,12 @@
2121
import com.google.api.gax.core.CredentialsProvider;
2222
import com.google.api.gax.core.GaxProperties;
2323
import com.google.api.gax.grpc.GaxGrpcProperties;
24+
import com.google.api.gax.grpc.GrpcCallContext;
2425
import com.google.api.gax.grpc.GrpcTransportChannel;
2526
import com.google.api.gax.rpc.ApiClientHeaderProvider;
2627
import com.google.api.gax.rpc.FixedTransportChannelProvider;
2728
import com.google.api.gax.rpc.HeaderProvider;
29+
import com.google.api.gax.rpc.ServerStream;
2830
import com.google.api.gax.rpc.TransportChannelProvider;
2931
import com.google.api.pathtemplate.PathTemplate;
3032
import com.google.cloud.ServiceOptions;
@@ -72,6 +74,7 @@
7274
import com.google.spanner.v1.PartitionQueryRequest;
7375
import com.google.spanner.v1.PartitionReadRequest;
7476
import com.google.spanner.v1.PartitionResponse;
77+
import com.google.spanner.v1.PartialResultSet;
7578
import com.google.spanner.v1.ReadRequest;
7679
import com.google.spanner.v1.RollbackRequest;
7780
import com.google.spanner.v1.Session;
@@ -335,15 +338,19 @@ public void deleteSession(String sessionName, @Nullable Map<Option, ?> options)
335338
}
336339

337340
@Override
338-
public StreamingCall read(
341+
public ServerStream<PartialResultSet> read(
339342
ReadRequest request, ResultStreamConsumer consumer, @Nullable Map<Option, ?> options) {
340-
throw new UnsupportedOperationException("Not implemented yet.");
343+
GrpcCallContext context = GrpcCallContext.createDefault()
344+
.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue());
345+
return stub.streamingReadCallable().call(request, context);
341346
}
342347

343348
@Override
344-
public StreamingCall executeQuery(
349+
public ServerStream<PartialResultSet> executeQuery(
345350
ExecuteSqlRequest request, ResultStreamConsumer consumer, @Nullable Map<Option, ?> options) {
346-
throw new UnsupportedOperationException("Not implemented yet.");
351+
GrpcCallContext context = GrpcCallContext.createDefault()
352+
.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue());
353+
return stub.executeStreamingSqlCallable().call(request, context);
347354
}
348355

349356
@Override

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GrpcSpannerRpc.java‎

Lines changed: 5 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import com.google.api.gax.grpc.GaxGrpcProperties;
2323
import com.google.api.gax.rpc.ApiClientHeaderProvider;
2424
import com.google.api.gax.rpc.HeaderProvider;
25+
import com.google.api.gax.rpc.ServerStream;
2526
import com.google.api.pathtemplate.PathTemplate;
2627
import com.google.cloud.NoCredentials;
2728
import com.google.cloud.ServiceOptions;
@@ -366,25 +367,15 @@ public void deleteSession(String sessionName, @Nullable Map<Option, ?> options)
366367
}
367368

368369
@Override
369-
public StreamingCall read(
370+
public ServerStream<PartialResultSet> read(
370371
ReadRequest request, ResultStreamConsumer consumer, @Nullable Map<Option, ?> options) {
371-
return doStreamingCall(
372-
SpannerGrpc.METHOD_STREAMING_READ,
373-
request,
374-
consumer,
375-
request.getSession(),
376-
Option.CHANNEL_HINT.getLong(options));
372+
throw new UnsupportedOperationException("Not implemented: read");
377373
}
378374

379375
@Override
380-
public StreamingCall executeQuery(
376+
public ServerStream<PartialResultSet> executeQuery(
381377
ExecuteSqlRequest request, ResultStreamConsumer consumer, @Nullable Map<Option, ?> options) {
382-
return doStreamingCall(
383-
SpannerGrpc.METHOD_EXECUTE_STREAMING_SQL,
384-
request,
385-
consumer,
386-
request.getSession(),
387-
Option.CHANNEL_HINT.getLong(options));
378+
throw new UnsupportedOperationException("Not implemented: executeQuery");
388379
}
389380

390381
@Override

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
package com.google.cloud.spanner.spi.v1;
1818

19+
import com.google.api.gax.rpc.ServerStream;
1920
import com.google.cloud.ServiceRpc;
2021
import com.google.cloud.spanner.SpannerException;
2122
import com.google.cloud.spanner.spi.v1.SpannerRpc.Option;
@@ -197,10 +198,10 @@ Session createSession(String databaseName, @Nullable Map<String, String> labels,
197198

198199
void deleteSession(String sessionName, @Nullable Map<Option, ?> options) throws SpannerException;
199200

200-
StreamingCall read(
201+
ServerStream<PartialResultSet> read(
201202
ReadRequest request, ResultStreamConsumer consumer, @Nullable Map<Option, ?> options);
202203

203-
StreamingCall executeQuery(
204+
ServerStream<PartialResultSet> executeQuery(
204205
ExecuteSqlRequest request, ResultStreamConsumer consumer, @Nullable Map<Option, ?> options);
205206

206207
Transaction beginTransaction(BeginTransactionRequest request, @Nullable Map<Option, ?> options)

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/ServerStreamingStashCallable.java‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
/*
2+
* Copyright 2018 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 com.google.api.gax.rpc.ApiCallContext;
20+
import com.google.api.gax.rpc.ResponseObserver;
21+
import com.google.api.gax.rpc.ServerStream;
22+
import com.google.api.gax.rpc.ServerStreamingCallable;
23+
import com.google.api.gax.rpc.StreamController;
24+
import com.google.common.base.Preconditions;
25+
import com.google.common.collect.Queues;
26+
import java.util.ArrayList;
27+
import java.util.List;
28+
import java.util.Queue;
29+
import java.util.concurrent.CancellationException;
30+
31+
public class ServerStreamingStashCallable<RequestT, ResponseT>
32+
extends ServerStreamingCallable<RequestT, ResponseT> {
33+
private List<ResponseT> responseList;
34+
35+
public ServerStreamingStashCallable() {
36+
responseList = new ArrayList<>();
37+
}
38+
39+
public ServerStreamingStashCallable(List<ResponseT> responseList) {
40+
this.responseList = responseList;
41+
}
42+
43+
@Override
44+
public void call(
45+
RequestT request, ResponseObserver<ResponseT> responseObserver, ApiCallContext context) {
46+
Preconditions.checkNotNull(responseObserver);
47+
48+
StreamControllerStash<ResponseT> controller =
49+
new StreamControllerStash<>(responseList, responseObserver);
50+
controller.start();
51+
}
52+
53+
// Minimal implementation of back pressure aware stream controller. Not threadsafe
54+
private static class StreamControllerStash<ResponseT> implements StreamController {
55+
final ResponseObserver<ResponseT> observer;
56+
final Queue<ResponseT> queue;
57+
boolean autoFlowControl = true;
58+
long numPending;
59+
Throwable error;
60+
boolean delivering, closed;
61+
62+
public StreamControllerStash(
63+
List<ResponseT> responseList, ResponseObserver<ResponseT> observer) {
64+
this.observer = observer;
65+
this.queue = Queues.newArrayDeque(responseList);
66+
}
67+
68+
public void start() {
69+
observer.onStart(this);
70+
if (autoFlowControl) {
71+
numPending = Integer.MAX_VALUE;
72+
}
73+
deliver();
74+
}
75+
76+
@Override
77+
public void disableAutoInboundFlowControl() {
78+
autoFlowControl = false;
79+
}
80+
81+
@Override
82+
public void request(int count) {
83+
numPending += count;
84+
deliver();
85+
}
86+
87+
@Override
88+
public void cancel() {
89+
error = new CancellationException("User cancelled stream");
90+
deliver();
91+
}
92+
93+
private void deliver() {
94+
if (delivering || closed) return;
95+
delivering = true;
96+
97+
try {
98+
while (error == null && numPending > 0 && !queue.isEmpty()) {
99+
numPending--;
100+
observer.onResponse(queue.poll());
101+
}
102+
103+
if (error != null || queue.isEmpty()) {
104+
if (error != null) {
105+
observer.onError(error);
106+
} else {
107+
observer.onComplete();
108+
}
109+
closed = true;
110+
}
111+
} finally {
112+
delivering = false;
113+
}
114+
}
115+
}
116+
}

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java‎

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
import static com.google.common.truth.Truth.assertThat;
2020
import static org.junit.Assert.fail;
2121

22+
import com.google.api.gax.rpc.ServerStream;
23+
import com.google.api.gax.rpc.ServerStreamingCallable;
2224
import com.google.cloud.Timestamp;
2325
import com.google.cloud.spanner.spi.v1.SpannerRpc;
2426
import com.google.protobuf.ByteString;
@@ -280,18 +282,15 @@ public void request(int numMessages) {}
280282
}
281283

282284
private void mockRead(final PartialResultSet myResultSet) {
283-
final ArgumentCaptor<SpannerRpc.ResultStreamConsumer> consumer =
284-
ArgumentCaptor.forClass(SpannerRpc.ResultStreamConsumer.class);
285-
Mockito.when(rpc.read(Mockito.<ReadRequest>any(), consumer.capture(), Mockito.eq(options)))
286-
.then(
287-
new Answer<SpannerRpc.StreamingCall>() {
288-
@Override
289-
public SpannerRpc.StreamingCall answer(InvocationOnMock invocation) throws Throwable {
290-
consumer.getValue().onPartialResultSet(myResultSet);
291-
consumer.getValue().onCompleted();
292-
return new NoOpStreamingCall();
293-
}
294-
});
285+
ServerStreamingCallable<ReadRequest, PartialResultSet> serverStreamingCallable =
286+
new ServerStreamingStashCallable(Arrays.<PartialResultSet>asList(myResultSet));
287+
final ServerStream<PartialResultSet> mockServerStream = serverStreamingCallable.call(null);
288+
Mockito.when(
289+
rpc.read(
290+
Mockito.<ReadRequest>any(),
291+
Mockito.<SpannerRpc.ResultStreamConsumer>any(),
292+
Mockito.eq(options)))
293+
.thenReturn(mockServerStream);
295294
}
296295

297296
@Test

0 commit comments

Comments
 (0)