Skip to content
This repository was archived by the owner on Apr 7, 2026. It is now read-only.

Commit 33175a5

Browse files
rahul2393rahul2393
committed
add more tests
1 parent d719633 commit 33175a5

5 files changed

Lines changed: 594 additions & 64 deletions

File tree

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

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -399,9 +399,7 @@ public GapicSpannerRpc(final SpannerOptions options) {
399399
// If it is enabled in options uses the channel pool provided by the gRPC-GCP extension.
400400
maybeEnableGrpcGcpExtension(defaultChannelProviderBuilder, options);
401401

402-
boolean enableLocationApi =
403-
options.isEnableLocationApi()
404-
|| Boolean.parseBoolean(System.getenv(EXPERIMENTAL_LOCATION_API_ENV_VAR));
402+
boolean enableLocationApi = options.isEnableLocationApi();
405403
TransportChannelProvider baseChannelProvider =
406404
MoreObjects.firstNonNull(
407405
options.getChannelProvider(), defaultChannelProviderBuilder.build());

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

Lines changed: 13 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -248,9 +248,8 @@ static final class KeyAwareClientCall<RequestT, ResponseT>
248248
private long pendingRequests;
249249
private boolean pendingHalfClose;
250250
@Nullable private Boolean pendingMessageCompression;
251-
private boolean cancelled;
252-
@Nullable private String cancelMessage;
253-
@Nullable private Throwable cancelCause;
251+
@Nullable private io.grpc.Status cancelledStatus;
252+
@Nullable private Metadata cancelledTrailers;
254253

255254
KeyAwareClientCall(
256255
KeyAwareChannel parentChannel,
@@ -274,17 +273,17 @@ protected ClientCall<RequestT, ResponseT> delegate() {
274273
public void start(Listener<ResponseT> responseListener, Metadata headers) {
275274
this.responseListener = new KeyAwareClientCallListener<>(responseListener, this);
276275
this.headers = headers;
277-
if (cancelled) {
276+
if (this.cancelledStatus != null) {
278277
this.responseListener.onClose(
279-
io.grpc.Status.CANCELLED.withDescription(cancelMessage).withCause(cancelCause),
280-
new Metadata());
278+
this.cancelledStatus,
279+
this.cancelledTrailers == null ? new Metadata() : this.cancelledTrailers);
281280
}
282281
}
283282

284283
@Override
285284
@SuppressWarnings("unchecked")
286285
public void sendMessage(RequestT message) {
287-
if (cancelled) {
286+
if (this.cancelledStatus != null) {
288287
return;
289288
}
290289
if (responseListener == null || headers == null) {
@@ -311,10 +310,7 @@ public void sendMessage(RequestT message) {
311310
String databaseId = parentChannel.extractDatabaseIdFromSession(reqBuilder.getSession());
312311
if (databaseId != null && reqBuilder.hasMutationKey()) {
313312
finder = parentChannel.getOrCreateChannelFinder(databaseId);
314-
ChannelEndpoint routed = finder.findServer(reqBuilder);
315-
if (endpoint == null) {
316-
endpoint = routed;
317-
}
313+
endpoint = finder.findServer(reqBuilder);
318314
}
319315
allowDefaultAffinity = true;
320316
message = (RequestT) reqBuilder.build();
@@ -345,6 +341,7 @@ public void sendMessage(RequestT message) {
345341
delegate = endpoint.getChannel().newCall(methodDescriptor, callOptions);
346342
if (pendingMessageCompression != null) {
347343
delegate.setMessageCompression(pendingMessageCompression);
344+
pendingMessageCompression = null;
348345
}
349346
delegate.start(responseListener, headers);
350347
drainPendingRequests();
@@ -368,12 +365,12 @@ public void cancel(@Nullable String message, @Nullable Throwable cause) {
368365
if (delegate != null) {
369366
delegate.cancel(message, cause);
370367
} else {
371-
cancelled = true;
372-
cancelMessage = message;
373-
cancelCause = cause;
368+
cancelledStatus = io.grpc.Status.CANCELLED.withDescription(message).withCause(cause);
369+
Metadata trailers =
370+
cause == null ? new Metadata() : io.grpc.Status.trailersFromThrowable(cause);
371+
cancelledTrailers = trailers == null ? new Metadata() : trailers;
374372
if (responseListener != null) {
375-
responseListener.onClose(
376-
io.grpc.Status.CANCELLED.withDescription(message).withCause(cause), new Metadata());
373+
responseListener.onClose(cancelledStatus, cancelledTrailers);
377374
}
378375
}
379376
}

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

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,21 +16,39 @@
1616

1717
package com.google.cloud.spanner;
1818

19+
import static com.google.cloud.spanner.SpannerApiFutures.get;
1920
import static org.junit.Assert.assertEquals;
21+
import static org.junit.Assert.assertThrows;
22+
import static org.junit.Assert.assertTrue;
2023

24+
import com.google.api.core.ApiFuture;
25+
import com.google.api.gax.rpc.ApiCallContext;
2126
import com.google.cloud.NoCredentials;
27+
import com.google.cloud.spanner.AsyncResultSet.CallbackResponse;
28+
import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
2229
import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
30+
import com.google.cloud.spanner.SpannerOptions.CallContextConfigurator;
2331
import com.google.cloud.spanner.connection.AbstractMockServerTest;
2432
import com.google.cloud.spanner.connection.RandomResultSetGenerator;
2533
import com.google.common.util.concurrent.Futures;
2634
import com.google.common.util.concurrent.ListenableFuture;
2735
import com.google.common.util.concurrent.ListeningExecutorService;
2836
import com.google.common.util.concurrent.MoreExecutors;
37+
import com.google.spanner.v1.ExecuteSqlRequest;
38+
import com.google.spanner.v1.SpannerGrpc;
39+
import io.grpc.Context;
2940
import io.grpc.ManagedChannelBuilder;
41+
import io.grpc.MethodDescriptor;
42+
import io.grpc.Status;
43+
import java.time.Duration;
3044
import java.util.ArrayList;
45+
import java.util.LinkedList;
3146
import java.util.List;
47+
import java.util.concurrent.CountDownLatch;
48+
import java.util.concurrent.ExecutorService;
3249
import java.util.concurrent.Executors;
3350
import java.util.concurrent.ThreadLocalRandom;
51+
import java.util.concurrent.TimeUnit;
3452
import org.junit.AfterClass;
3553
import org.junit.BeforeClass;
3654
import org.junit.Test;
@@ -44,6 +62,10 @@ public class LocationAwareTest extends AbstractMockServerTest {
4462
private static Spanner spanner;
4563
private static DatabaseClient client;
4664

65+
private static final class TimeoutHolder {
66+
private Duration timeout;
67+
}
68+
4769
@BeforeClass
4870
public static void enableLocationApiAndSetupClient() {
4971
SpannerOptions.useEnvironment(
@@ -137,4 +159,112 @@ public void testParallelReadWriteTransactions() throws Exception {
137159
executor.shutdown();
138160
Futures.allAsList(results).get();
139161
}
162+
163+
@Test
164+
public void testExecuteStreamingSqlCallContextTimeout_locationAware() {
165+
final TimeoutHolder timeoutHolder = new TimeoutHolder();
166+
CallContextConfigurator configurator =
167+
new CallContextConfigurator() {
168+
@Override
169+
public <ReqT, RespT> ApiCallContext configure(
170+
ApiCallContext context, ReqT request, MethodDescriptor<ReqT, RespT> method) {
171+
if (request instanceof ExecuteSqlRequest
172+
&& method.equals(SpannerGrpc.getExecuteStreamingSqlMethod())) {
173+
return context.withTimeoutDuration(timeoutHolder.timeout);
174+
}
175+
return null;
176+
}
177+
};
178+
179+
mockSpanner.setExecuteStreamingSqlExecutionTime(
180+
SimulatedExecutionTime.ofMinimumAndRandomTime(10, 0));
181+
Context context =
182+
Context.current().withValue(SpannerOptions.CALL_CONTEXT_CONFIGURATOR_KEY, configurator);
183+
try {
184+
context.run(
185+
() -> {
186+
timeoutHolder.timeout = Duration.ofNanos(1L);
187+
SpannerException e =
188+
assertThrows(
189+
SpannerException.class,
190+
() -> {
191+
try (ResultSet rs =
192+
client.singleUse().executeQuery(SELECT_RANDOM_STATEMENT)) {
193+
rs.next();
194+
}
195+
});
196+
assertEquals(ErrorCode.DEADLINE_EXCEEDED, e.getErrorCode());
197+
198+
timeoutHolder.timeout = Duration.ofMinutes(1L);
199+
try (ResultSet rs = client.singleUse().executeQuery(SELECT_RANDOM_STATEMENT)) {
200+
assertTrue(rs.next());
201+
}
202+
});
203+
} finally {
204+
mockSpanner.removeAllExecutionTimes();
205+
}
206+
}
207+
208+
@Test
209+
public void testExecuteStreamingSqlInvalidArgumentPropagates_locationAware() {
210+
mockSpanner.setExecuteStreamingSqlExecutionTime(
211+
SimulatedExecutionTime.ofException(
212+
Status.INVALID_ARGUMENT.withDescription("invalid request").asRuntimeException()));
213+
try {
214+
SpannerException e =
215+
assertThrows(
216+
SpannerException.class,
217+
() -> {
218+
try (ResultSet rs = client.singleUse().executeQuery(SELECT_RANDOM_STATEMENT)) {
219+
rs.next();
220+
}
221+
});
222+
assertEquals(ErrorCode.INVALID_ARGUMENT, e.getErrorCode());
223+
} finally {
224+
mockSpanner.removeAllExecutionTimes();
225+
}
226+
}
227+
228+
@Test
229+
public void testExecuteQueryAsyncCancelReturnsCancelled_locationAware() throws Exception {
230+
final List<Integer> values = new LinkedList<>();
231+
final CountDownLatch receivedFirstRow = new CountDownLatch(1);
232+
final CountDownLatch cancelled = new CountDownLatch(1);
233+
final ApiFuture<Void> callbackResult;
234+
235+
ExecutorService executor = Executors.newSingleThreadExecutor();
236+
try (AsyncResultSet rs = client.singleUse().executeQueryAsync(SELECT_RANDOM_STATEMENT)) {
237+
callbackResult =
238+
rs.setCallback(
239+
executor,
240+
resultSet -> {
241+
try {
242+
while (true) {
243+
switch (resultSet.tryNext()) {
244+
case DONE:
245+
return CallbackResponse.DONE;
246+
case NOT_READY:
247+
return CallbackResponse.CONTINUE;
248+
case OK:
249+
values.add(1);
250+
receivedFirstRow.countDown();
251+
cancelled.await();
252+
break;
253+
}
254+
}
255+
} catch (Throwable t) {
256+
return CallbackResponse.DONE;
257+
}
258+
});
259+
260+
assertTrue(receivedFirstRow.await(30L, TimeUnit.SECONDS));
261+
rs.cancel();
262+
cancelled.countDown();
263+
SpannerException e = assertThrows(SpannerException.class, () -> get(callbackResult));
264+
assertEquals(ErrorCode.CANCELLED, e.getErrorCode());
265+
assertEquals(1, values.size());
266+
} finally {
267+
executor.shutdownNow();
268+
}
269+
}
140270
}

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java‎

Lines changed: 45 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,6 @@
4646
import com.google.cloud.spanner.DatabaseId;
4747
import com.google.cloud.spanner.Dialect;
4848
import com.google.cloud.spanner.ErrorCode;
49-
import com.google.cloud.spanner.JavaVersionUtil;
5049
import com.google.cloud.spanner.MockSpannerServiceImpl;
5150
import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
5251
import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
@@ -888,18 +887,7 @@ public void testCreateSession_whenMultiplexedSessionIsFalse_assertSessionProto()
888887
}
889888

890889
@Test
891-
public void testChannelEndpointCacheFactoryUsedWhenLocationApiEnabled() throws Exception {
892-
assumeTrue(isJava8() && !isWindows());
893-
String envVar = "GOOGLE_SPANNER_EXPERIMENTAL_LOCATION_API";
894-
895-
Class<?> classOfMap = System.getenv().getClass();
896-
java.lang.reflect.Field field = classOfMap.getDeclaredField("m");
897-
field.setAccessible(true);
898-
@SuppressWarnings("unchecked")
899-
Map<String, String> writeableEnvironmentVariables =
900-
(Map<String, String>) field.get(System.getenv());
901-
String originalValue = writeableEnvironmentVariables.get(envVar);
902-
890+
public void testChannelEndpointCacheFactoryUsedWhenLocationApiEnabled() {
903891
AtomicBoolean factoryCalled = new AtomicBoolean(false);
904892
ChannelEndpointCacheFactory factory =
905893
baseProvider -> {
@@ -908,34 +896,25 @@ public void testChannelEndpointCacheFactoryUsedWhenLocationApiEnabled() throws E
908896
};
909897

910898
try {
911-
writeableEnvironmentVariables.put(envVar, "true");
899+
SpannerOptions.useEnvironment(
900+
new SpannerOptions.SpannerEnvironment() {
901+
@Override
902+
public boolean isEnableLocationApi() {
903+
return true;
904+
}
905+
});
912906
SpannerOptions options =
913907
createSpannerOptions().toBuilder().setChannelEndpointCacheFactory(factory).build();
914908
GapicSpannerRpc rpc = new GapicSpannerRpc(options, true);
915909
rpc.shutdown();
916910
assertTrue(factoryCalled.get());
917911
} finally {
918-
if (originalValue == null) {
919-
writeableEnvironmentVariables.remove(envVar);
920-
} else {
921-
writeableEnvironmentVariables.put(envVar, originalValue);
922-
}
912+
SpannerOptions.useDefaultEnvironment();
923913
}
924914
}
925915

926916
@Test
927-
public void testLocationApiDoesNotOverrideExplicitChannelProvider() throws Exception {
928-
assumeTrue(isJava8() && !isWindows());
929-
String envVar = "GOOGLE_SPANNER_EXPERIMENTAL_LOCATION_API";
930-
931-
Class<?> classOfMap = System.getenv().getClass();
932-
java.lang.reflect.Field field = classOfMap.getDeclaredField("m");
933-
field.setAccessible(true);
934-
@SuppressWarnings("unchecked")
935-
Map<String, String> writeableEnvironmentVariables =
936-
(Map<String, String>) field.get(System.getenv());
937-
String originalValue = writeableEnvironmentVariables.get(envVar);
938-
917+
public void testLocationApiDoesNotOverrideExplicitChannelProvider() {
939918
AtomicBoolean factoryCalled = new AtomicBoolean(false);
940919
ChannelEndpointCacheFactory factory =
941920
baseProvider -> {
@@ -949,7 +928,13 @@ public void testLocationApiDoesNotOverrideExplicitChannelProvider() throws Excep
949928
address.getHostString(), server.getPort(), providerUsed);
950929

951930
try {
952-
writeableEnvironmentVariables.put(envVar, "true");
931+
SpannerOptions.useEnvironment(
932+
new SpannerOptions.SpannerEnvironment() {
933+
@Override
934+
public boolean isEnableLocationApi() {
935+
return true;
936+
}
937+
});
953938
SpannerOptions options =
954939
createSpannerOptions().toBuilder()
955940
.setChannelProvider(channelProvider)
@@ -960,11 +945,34 @@ public void testLocationApiDoesNotOverrideExplicitChannelProvider() throws Excep
960945
assertTrue(providerUsed.get());
961946
assertFalse(factoryCalled.get());
962947
} finally {
963-
if (originalValue == null) {
964-
writeableEnvironmentVariables.remove(envVar);
965-
} else {
966-
writeableEnvironmentVariables.put(envVar, originalValue);
967-
}
948+
SpannerOptions.useDefaultEnvironment();
949+
}
950+
}
951+
952+
@Test
953+
public void testLocationApiDisabledInOptionsDoesNotCreateKeyAwareChannelProvider() {
954+
AtomicBoolean factoryCalled = new AtomicBoolean(false);
955+
ChannelEndpointCacheFactory factory =
956+
baseProvider -> {
957+
factoryCalled.set(true);
958+
return new GrpcChannelEndpointCache(baseProvider);
959+
};
960+
961+
try {
962+
SpannerOptions.useEnvironment(
963+
new SpannerOptions.SpannerEnvironment() {
964+
@Override
965+
public boolean isEnableLocationApi() {
966+
return false;
967+
}
968+
});
969+
SpannerOptions options =
970+
createSpannerOptions().toBuilder().setChannelEndpointCacheFactory(factory).build();
971+
GapicSpannerRpc rpc = new GapicSpannerRpc(options, true);
972+
rpc.shutdown();
973+
assertFalse(factoryCalled.get());
974+
} finally {
975+
SpannerOptions.useDefaultEnvironment();
968976
}
969977
}
970978

@@ -1117,12 +1125,4 @@ private SpannerOptions createSpannerOptions() {
11171125
.setCallCredentialsProvider(() -> MoreCallCredentials.from(VARIABLE_CREDENTIALS))
11181126
.build();
11191127
}
1120-
1121-
private boolean isJava8() {
1122-
return JavaVersionUtil.getJavaMajorVersion() == 8;
1123-
}
1124-
1125-
private boolean isWindows() {
1126-
return System.getProperty("os.name").toLowerCase().contains("windows");
1127-
}
11281128
}

0 commit comments

Comments
 (0)