Skip to content

Commit 81f9ecc

Browse files
olavloiteolavloite
committed
examples: add more examples
1 parent c411f6c commit 81f9ecc

File tree

3 files changed

+283
-5
lines changed

3 files changed

+283
-5
lines changed

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package com.google.cloud.spanner;
1818

1919
import com.google.cloud.spanner.Type.StructField;
20+
import com.google.common.collect.ContiguousSet;
2021
import com.google.protobuf.ListValue;
2122
import com.google.spanner.v1.ResultSetMetadata;
2223
import com.google.spanner.v1.StructType;
@@ -108,11 +109,11 @@ public class MockSpannerTestUtil {
108109
.setMetadata(READ_KEY_VALUE_METADATA)
109110
.build();
110111
static final com.google.spanner.v1.ResultSet READ_MULTIPLE_KEY_VALUE_RESULTSET =
111-
generateKeyValueResultSet(1, 3);
112+
generateKeyValueResultSet(ContiguousSet.closed(1, 3));
112113

113-
static com.google.spanner.v1.ResultSet generateKeyValueResultSet(int beginRow, int endRow) {
114+
static com.google.spanner.v1.ResultSet generateKeyValueResultSet(Iterable<Integer> rows) {
114115
com.google.spanner.v1.ResultSet.Builder builder = com.google.spanner.v1.ResultSet.newBuilder();
115-
for (int row = beginRow; row <= endRow; row++) {
116+
for (Integer row : rows) {
116117
builder.addRows(
117118
ListValue.newBuilder()
118119
.addValues(com.google.protobuf.Value.newBuilder().setStringValue("k" + row).build())

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

Lines changed: 141 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,9 @@
3131
import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
3232
import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
3333
import com.google.common.base.Function;
34+
import com.google.common.collect.ContiguousSet;
3435
import com.google.common.collect.ImmutableList;
36+
import com.google.common.collect.ImmutableSet;
3537
import com.google.common.collect.Iterables;
3638
import com.google.common.util.concurrent.SettableFuture;
3739
import io.grpc.Server;
@@ -40,6 +42,7 @@
4042
import java.util.ArrayList;
4143
import java.util.Arrays;
4244
import java.util.Comparator;
45+
import java.util.LinkedList;
4346
import java.util.List;
4447
import java.util.concurrent.BlockingQueue;
4548
import java.util.concurrent.CountDownLatch;
@@ -282,9 +285,9 @@ public void readOnlyTransaction() throws Exception {
282285
Statement.of("SELECT * FROM TestTable WHERE Key IN ('k10', 'k11', 'k12')");
283286
Statement statement2 = Statement.of("SELECT * FROM TestTable WHERE Key IN ('k1', 'k2', 'k3");
284287
mockSpanner.putStatementResult(
285-
StatementResult.query(statement1, generateKeyValueResultSet(10, 12)));
288+
StatementResult.query(statement1, generateKeyValueResultSet(ContiguousSet.closed(10, 12))));
286289
mockSpanner.putStatementResult(
287-
StatementResult.query(statement2, generateKeyValueResultSet(1, 3)));
290+
StatementResult.query(statement2, generateKeyValueResultSet(ContiguousSet.closed(1, 3))));
288291

289292
ApiFuture<ImmutableList<String>> values1;
290293
ApiFuture<ImmutableList<String>> values2;
@@ -333,4 +336,140 @@ public int compare(String o1, String o2) {
333336
executor);
334337
assertThat(allValues.get()).containsExactly("v1", "v2", "v3", "v10", "v11", "v12");
335338
}
339+
340+
@Test
341+
public void pauseResume() throws Exception {
342+
Statement unevenStatement =
343+
Statement.of("SELECT * FROM TestTable WHERE MOD(CAST(SUBSTR(Key, 2) AS INT64), 2) = 1");
344+
Statement evenStatement =
345+
Statement.of("SELECT * FROM TestTable WHERE MOD(CAST(SUBSTR(Key, 2) AS INT64), 2) = 0");
346+
mockSpanner.putStatementResult(
347+
StatementResult.query(
348+
unevenStatement, generateKeyValueResultSet(ImmutableSet.of(1, 3, 5, 7, 9))));
349+
mockSpanner.putStatementResult(
350+
StatementResult.query(
351+
evenStatement, generateKeyValueResultSet(ImmutableSet.of(2, 4, 6, 8, 10))));
352+
353+
final SettableApiFuture<Boolean> evenFinished = SettableApiFuture.create();
354+
final SettableApiFuture<Boolean> unevenFinished = SettableApiFuture.create();
355+
final List<String> allValues = new LinkedList<>();
356+
try (ReadOnlyTransaction tx = client.readOnlyTransaction()) {
357+
try (AsyncResultSet evenRs = tx.executeQueryAsync(evenStatement);
358+
AsyncResultSet unevenRs = tx.executeQueryAsync(unevenStatement)) {
359+
evenRs.setCallback(
360+
executor,
361+
new ReadyCallback() {
362+
private boolean firstRow = true;
363+
364+
@Override
365+
public CallbackResponse cursorReady(AsyncResultSet resultSet) {
366+
if (firstRow) {
367+
// Make sure the uneven result set returns the first result.
368+
firstRow = false;
369+
return CallbackResponse.PAUSE;
370+
}
371+
try {
372+
while (true) {
373+
switch (resultSet.tryNext()) {
374+
case DONE:
375+
evenFinished.set(true);
376+
return CallbackResponse.DONE;
377+
case NOT_READY:
378+
return CallbackResponse.CONTINUE;
379+
case OK:
380+
allValues.add(resultSet.getString("Value"));
381+
return CallbackResponse.PAUSE;
382+
}
383+
}
384+
} catch (Throwable t) {
385+
evenFinished.setException(t);
386+
return CallbackResponse.DONE;
387+
} finally {
388+
unevenRs.resume();
389+
}
390+
}
391+
});
392+
393+
unevenRs.setCallback(
394+
executor,
395+
new ReadyCallback() {
396+
@Override
397+
public CallbackResponse cursorReady(AsyncResultSet resultSet) {
398+
try {
399+
while (true) {
400+
switch (resultSet.tryNext()) {
401+
case DONE:
402+
unevenFinished.set(true);
403+
return CallbackResponse.DONE;
404+
case NOT_READY:
405+
return CallbackResponse.CONTINUE;
406+
case OK:
407+
allValues.add(resultSet.getString("Value"));
408+
return CallbackResponse.PAUSE;
409+
}
410+
}
411+
} catch (Throwable t) {
412+
unevenFinished.setException(t);
413+
return CallbackResponse.DONE;
414+
} finally {
415+
evenRs.resume();
416+
}
417+
}
418+
});
419+
}
420+
}
421+
assertThat(ApiFutures.allAsList(Arrays.asList(evenFinished, unevenFinished)).get())
422+
.containsExactly(Boolean.TRUE, Boolean.TRUE);
423+
assertThat(allValues)
424+
.containsExactly("v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10");
425+
}
426+
427+
@Test
428+
public void cancel() throws Exception {
429+
final List<String> values = new LinkedList<>();
430+
final SettableApiFuture<Boolean> finished = SettableApiFuture.create();
431+
final CountDownLatch receivedFirstRow = new CountDownLatch(1);
432+
final CountDownLatch cancelled = new CountDownLatch(1);
433+
try (AsyncResultSet rs =
434+
client.singleUse().readAsync(READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES)) {
435+
rs.setCallback(
436+
executor,
437+
new ReadyCallback() {
438+
@Override
439+
public CallbackResponse cursorReady(AsyncResultSet resultSet) {
440+
try {
441+
while (true) {
442+
switch (resultSet.tryNext()) {
443+
case DONE:
444+
finished.set(true);
445+
return CallbackResponse.DONE;
446+
case NOT_READY:
447+
return CallbackResponse.CONTINUE;
448+
case OK:
449+
values.add(resultSet.getString("Value"));
450+
receivedFirstRow.countDown();
451+
cancelled.await();
452+
break;
453+
}
454+
}
455+
} catch (Throwable t) {
456+
finished.setException(t);
457+
return CallbackResponse.DONE;
458+
}
459+
}
460+
});
461+
receivedFirstRow.await();
462+
rs.cancel();
463+
}
464+
cancelled.countDown();
465+
try {
466+
finished.get();
467+
fail("missing expected exception");
468+
} catch (ExecutionException e) {
469+
assertThat(e.getCause()).isInstanceOf(SpannerException.class);
470+
SpannerException se = (SpannerException) e.getCause();
471+
assertThat(se.getErrorCode()).isEqualTo(ErrorCode.CANCELLED);
472+
assertThat(values).containsExactly("v1");
473+
}
474+
}
336475
}

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package com.google.cloud.spanner.it;
1818

1919
import static com.google.common.truth.Truth.assertThat;
20+
import static org.junit.Assert.fail;
2021

2122
import com.google.api.core.ApiFunction;
2223
import com.google.api.core.ApiFuture;
@@ -29,12 +30,14 @@
2930
import com.google.cloud.spanner.AsyncRunner.AsyncWork;
3031
import com.google.cloud.spanner.Database;
3132
import com.google.cloud.spanner.DatabaseClient;
33+
import com.google.cloud.spanner.ErrorCode;
3234
import com.google.cloud.spanner.IntegrationTest;
3335
import com.google.cloud.spanner.IntegrationTestEnv;
3436
import com.google.cloud.spanner.Key;
3537
import com.google.cloud.spanner.KeySet;
3638
import com.google.cloud.spanner.Mutation;
3739
import com.google.cloud.spanner.ReadOnlyTransaction;
40+
import com.google.cloud.spanner.SpannerException;
3841
import com.google.cloud.spanner.Statement;
3942
import com.google.cloud.spanner.Struct;
4043
import com.google.cloud.spanner.StructReader;
@@ -47,6 +50,8 @@
4750
import java.util.Comparator;
4851
import java.util.LinkedList;
4952
import java.util.List;
53+
import java.util.concurrent.CountDownLatch;
54+
import java.util.concurrent.ExecutionException;
5055
import java.util.concurrent.ExecutorService;
5156
import java.util.concurrent.Executors;
5257
import org.junit.AfterClass;
@@ -328,4 +333,137 @@ public int compare(String o1, String o2) {
328333
executor);
329334
assertThat(allValues.get()).containsExactly("v1", "v2", "v3", "v10", "v11", "v12");
330335
}
336+
337+
@Test
338+
public void pauseResume() throws Exception {
339+
Statement unevenStatement =
340+
Statement.of(
341+
"SELECT * FROM TestTable WHERE MOD(CAST(SUBSTR(Key, 2) AS INT64), 2) = 1 ORDER BY CAST(SUBSTR(Key, 2) AS INT64)");
342+
Statement evenStatement =
343+
Statement.of(
344+
"SELECT * FROM TestTable WHERE MOD(CAST(SUBSTR(Key, 2) AS INT64), 2) = 0 ORDER BY CAST(SUBSTR(Key, 2) AS INT64)");
345+
346+
final SettableApiFuture<Boolean> evenFinished = SettableApiFuture.create();
347+
final SettableApiFuture<Boolean> unevenFinished = SettableApiFuture.create();
348+
final List<String> allValues = new LinkedList<>();
349+
try (ReadOnlyTransaction tx = client.readOnlyTransaction()) {
350+
try (AsyncResultSet evenRs = tx.executeQueryAsync(evenStatement);
351+
AsyncResultSet unevenRs = tx.executeQueryAsync(unevenStatement)) {
352+
evenRs.setCallback(
353+
executor,
354+
new ReadyCallback() {
355+
@Override
356+
public CallbackResponse cursorReady(AsyncResultSet resultSet) {
357+
try {
358+
while (true) {
359+
switch (resultSet.tryNext()) {
360+
case DONE:
361+
evenFinished.set(true);
362+
return CallbackResponse.DONE;
363+
case NOT_READY:
364+
return CallbackResponse.CONTINUE;
365+
case OK:
366+
allValues.add(resultSet.getString("StringValue"));
367+
return CallbackResponse.PAUSE;
368+
}
369+
}
370+
} catch (Throwable t) {
371+
evenFinished.setException(t);
372+
return CallbackResponse.DONE;
373+
} finally {
374+
unevenRs.resume();
375+
}
376+
}
377+
});
378+
379+
unevenRs.setCallback(
380+
executor,
381+
new ReadyCallback() {
382+
private boolean firstRow = true;
383+
384+
@Override
385+
public CallbackResponse cursorReady(AsyncResultSet resultSet) {
386+
if (firstRow) {
387+
// Make sure the even result set returns the first result.
388+
firstRow = false;
389+
return CallbackResponse.PAUSE;
390+
}
391+
try {
392+
while (true) {
393+
switch (resultSet.tryNext()) {
394+
case DONE:
395+
unevenFinished.set(true);
396+
return CallbackResponse.DONE;
397+
case NOT_READY:
398+
return CallbackResponse.CONTINUE;
399+
case OK:
400+
allValues.add(resultSet.getString("StringValue"));
401+
return CallbackResponse.PAUSE;
402+
}
403+
}
404+
} catch (Throwable t) {
405+
unevenFinished.setException(t);
406+
return CallbackResponse.DONE;
407+
} finally {
408+
evenRs.resume();
409+
}
410+
}
411+
});
412+
}
413+
}
414+
assertThat(ApiFutures.allAsList(Arrays.asList(evenFinished, unevenFinished)).get())
415+
.containsExactly(Boolean.TRUE, Boolean.TRUE);
416+
assertThat(allValues)
417+
.containsExactly(
418+
"v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11", "v12", "v13",
419+
"v14");
420+
}
421+
422+
@Test
423+
public void cancel() throws Exception {
424+
final List<String> values = new LinkedList<>();
425+
final SettableApiFuture<Boolean> finished = SettableApiFuture.create();
426+
final CountDownLatch receivedFirstRow = new CountDownLatch(1);
427+
final CountDownLatch cancelled = new CountDownLatch(1);
428+
try (AsyncResultSet rs = client.singleUse().readAsync(TABLE_NAME, KeySet.all(), ALL_COLUMNS)) {
429+
rs.setCallback(
430+
executor,
431+
new ReadyCallback() {
432+
@Override
433+
public CallbackResponse cursorReady(AsyncResultSet resultSet) {
434+
try {
435+
while (true) {
436+
switch (resultSet.tryNext()) {
437+
case DONE:
438+
finished.set(true);
439+
return CallbackResponse.DONE;
440+
case NOT_READY:
441+
return CallbackResponse.CONTINUE;
442+
case OK:
443+
values.add(resultSet.getString("StringValue"));
444+
receivedFirstRow.countDown();
445+
cancelled.await();
446+
break;
447+
}
448+
}
449+
} catch (Throwable t) {
450+
finished.setException(t);
451+
return CallbackResponse.DONE;
452+
}
453+
}
454+
});
455+
receivedFirstRow.await();
456+
rs.cancel();
457+
}
458+
cancelled.countDown();
459+
try {
460+
finished.get();
461+
fail("missing expected exception");
462+
} catch (ExecutionException e) {
463+
assertThat(e.getCause()).isInstanceOf(SpannerException.class);
464+
SpannerException se = (SpannerException) e.getCause();
465+
assertThat(se.getErrorCode()).isEqualTo(ErrorCode.CANCELLED);
466+
assertThat(values).containsExactly("v0");
467+
}
468+
}
331469
}

0 commit comments

Comments
 (0)