Examples with CallbackResponse com.google.cloud.spanner.AsyncResultSet.CallbackResponse used on opensource projects

Example 1 with CallbackResponse

use of com.google.cloud.spanner.AsyncResultSet.CallbackResponse in project java-spanner by googleapis.

the class ConnectionAsyncApiAbortedTest method testUpdateAndQueryAbortedMidway.

@Test
public void testUpdateAndQueryAbortedMidway() throws InterruptedException {
    mockSpanner.setExecuteStreamingSqlExecutionTime(SimulatedExecutionTime.ofStreamException(mockSpanner.createAbortedException(ByteString.copyFromUtf8("test")), RANDOM_RESULT_SET_ROW_COUNT / 2));
    final RetryCounter counter = new RetryCounter();
    try (Connection connection = createConnection(counter)) {
        assertThat(counter.retryCount).isEqualTo(0);
        final SettableApiFuture<Long> rowCount = SettableApiFuture.create();
        final CountDownLatch updateLatch = new CountDownLatch(1);
        final CountDownLatch queryLatch = new CountDownLatch(1);
        ApiFuture<Void> finished;
        try (AsyncResultSet rs = connection.executeQueryAsync(SELECT_RANDOM_STATEMENT, Options.bufferRows(RANDOM_RESULT_SET_ROW_COUNT / 2 - 1))) {
            finished = rs.setCallback(singleThreadedExecutor, new ReadyCallback() {

                long count;

                @Override
                public CallbackResponse cursorReady(AsyncResultSet resultSet) {
                    // Indicate that the query has been executed.
                    queryLatch.countDown();
                    try {
                        // Wait until the update is on its way.
                        updateLatch.await(10L, TimeUnit.SECONDS);
                        while (true) {
                            switch(resultSet.tryNext()) {
                                case OK:
                                    count++;
                                    break;
                                case DONE:
                                    rowCount.set(count);
                                    return CallbackResponse.DONE;
                                case NOT_READY:
                                    return CallbackResponse.CONTINUE;
                            }
                        }
                    } catch (InterruptedException e) {
                        throw SpannerExceptionFactory.propagateInterrupt(e);
                    }
                }
            });
        }
        // Wait until the query has actually executed.
        queryLatch.await(10L, TimeUnit.SECONDS);
        ApiFuture<Long> updateCount = connection.executeUpdateAsync(INSERT_STATEMENT);
        updateCount.addListener(updateLatch::countDown, MoreExecutors.directExecutor());
        // We should not commit before the AsyncResultSet has finished.
        assertThat(get(finished)).isNull();
        ApiFuture<Void> commit = connection.commitAsync();
        assertThat(get(rowCount)).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT);
        assertThat(get(updateCount)).isEqualTo(UPDATE_COUNT);
        assertThat(get(commit)).isNull();
        assertThat(counter.retryCount).isEqualTo(1);
        // Verify the order of the statements on the server.
        List<? extends AbstractMessage> requests = Lists.newArrayList(Collections2.filter(mockSpanner.getRequests(), input -> input instanceof ExecuteSqlRequest));
        // The entire transaction should be retried.
        assertThat(requests).hasSize(4);
        assertThat(((ExecuteSqlRequest) requests.get(0)).getSeqno()).isEqualTo(1L);
        assertThat(((ExecuteSqlRequest) requests.get(0)).getSql()).isEqualTo(SELECT_RANDOM_STATEMENT.getSql());
        assertThat(((ExecuteSqlRequest) requests.get(1)).getSeqno()).isEqualTo(2L);
        assertThat(((ExecuteSqlRequest) requests.get(1)).getSql()).isEqualTo(INSERT_STATEMENT.getSql());
        assertThat(((ExecuteSqlRequest) requests.get(2)).getSeqno()).isEqualTo(1L);
        assertThat(((ExecuteSqlRequest) requests.get(2)).getSql()).isEqualTo(SELECT_RANDOM_STATEMENT.getSql());
        assertThat(((ExecuteSqlRequest) requests.get(3)).getSeqno()).isEqualTo(2L);
        assertThat(((ExecuteSqlRequest) requests.get(3)).getSql()).isEqualTo(INSERT_STATEMENT.getSql());
    }
}

Example 2 with CallbackResponse

use of com.google.cloud.spanner.AsyncResultSet.CallbackResponse in project java-spanner by googleapis.

the class ConnectionAsyncApiAbortedTest method testQueriesAbortedMidway_ResultsChanged.

@Test
public void testQueriesAbortedMidway_ResultsChanged() {
    mockSpanner.setExecuteStreamingSqlExecutionTime(SimulatedExecutionTime.ofStreamException(mockSpanner.createAbortedException(ByteString.copyFromUtf8("test")), RANDOM_RESULT_SET_ROW_COUNT - 1));
    final Statement statement = Statement.of("SELECT * FROM TEST_TABLE");
    final RandomResultSetGenerator generator = new RandomResultSetGenerator(RANDOM_RESULT_SET_ROW_COUNT - 10);
    mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate()));
    final CountDownLatch latch = new CountDownLatch(1);
    final RetryCounter counter = new RetryCounter();
    try (Connection connection = createConnection(counter)) {
        ApiFuture<Void> res1;
        try (AsyncResultSet rs = connection.executeQueryAsync(SELECT_RANDOM_STATEMENT, Options.bufferRows(5))) {
            res1 = rs.setCallback(multiThreadedExecutor, resultSet -> {
                try {
                    latch.await(10L, TimeUnit.SECONDS);
                    while (true) {
                        switch(resultSet.tryNext()) {
                            case OK:
                                break;
                            case DONE:
                                return CallbackResponse.DONE;
                            case NOT_READY:
                                return CallbackResponse.CONTINUE;
                        }
                    }
                } catch (Throwable t) {
                    throw SpannerExceptionFactory.asSpannerException(t);
                }
            });
        }
        try (AsyncResultSet rs = connection.executeQueryAsync(statement, Options.bufferRows(5))) {
            rs.setCallback(multiThreadedExecutor, new ReadyCallback() {

                boolean replaced;

                @Override
                public CallbackResponse cursorReady(AsyncResultSet resultSet) {
                    if (!replaced) {
                        // Replace the result of the query on the server after the first execution.
                        mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate()));
                        replaced = true;
                    }
                    while (true) {
                        switch(resultSet.tryNext()) {
                            case OK:
                                break;
                            case DONE:
                                latch.countDown();
                                return CallbackResponse.DONE;
                            case NOT_READY:
                                return CallbackResponse.CONTINUE;
                        }
                    }
                }
            });
        }
        try {
            get(res1);
            fail("Missing expected exception");
        } catch (AbortedDueToConcurrentModificationException e) {
            assertThat(counter.retryCount).isEqualTo(1);
        }
    }
}

Example 3 with CallbackResponse

use of com.google.cloud.spanner.AsyncResultSet.CallbackResponse in project java-spanner by googleapis.

the class ConnectionAsyncApiTest method testReadWriteMultipleAsyncStatements.

@Test
public void testReadWriteMultipleAsyncStatements() {
    try (Connection connection = createConnection()) {
        assertThat(connection.isAutocommit()).isFalse();
        ApiFuture<Long> update1 = connection.executeUpdateAsync(INSERT_STATEMENT);
        ApiFuture<Long> update2 = connection.executeUpdateAsync(INSERT_STATEMENT);
        ApiFuture<long[]> batch = connection.executeBatchUpdateAsync(ImmutableList.of(INSERT_STATEMENT, INSERT_STATEMENT));
        final SettableApiFuture<Integer> rowCount = SettableApiFuture.create();
        try (AsyncResultSet rs = connection.executeQueryAsync(SELECT_RANDOM_STATEMENT)) {
            rs.setCallback(executor, new ReadyCallback() {

                int count = 0;

                @Override
                public CallbackResponse cursorReady(AsyncResultSet resultSet) {
                    try {
                        while (true) {
                            switch(resultSet.tryNext()) {
                                case DONE:
                                    rowCount.set(count);
                                    return CallbackResponse.DONE;
                                case NOT_READY:
                                    return CallbackResponse.CONTINUE;
                                case OK:
                                    count++;
                            }
                        }
                    } catch (SpannerException e) {
                        rowCount.setException(e);
                        return CallbackResponse.DONE;
                    }
                }
            });
        }
        ApiFuture<Void> commit = connection.commitAsync();
        assertThat(get(update1)).isEqualTo(UPDATE_COUNT);
        assertThat(get(update2)).isEqualTo(UPDATE_COUNT);
        assertThat(get(batch)).asList().containsExactly(1L, 1L);
        assertThat(get(rowCount)).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT);
        assertNull(get(commit));
        // Get the last commit request.
        CommitRequest commitRequest = mockSpanner.getRequestsOfType(CommitRequest.class).stream().reduce((first, second) -> second).get();
        // Verify the order of the statements on the server.
        List<? extends AbstractMessage> requests = Lists.newArrayList(Collections2.filter(mockSpanner.getRequests(), input -> (input instanceof ExecuteSqlRequest && ((ExecuteSqlRequest) input).getSession().equals(commitRequest.getSession())) || (input instanceof ExecuteBatchDmlRequest && ((ExecuteBatchDmlRequest) input).getSession().equals(commitRequest.getSession()))));
        assertThat(requests).hasSize(4);
        assertThat(requests.get(0)).isInstanceOf(ExecuteSqlRequest.class);
        assertThat(((ExecuteSqlRequest) requests.get(0)).getSeqno()).isEqualTo(1L);
        assertThat(requests.get(1)).isInstanceOf(ExecuteSqlRequest.class);
        assertThat(((ExecuteSqlRequest) requests.get(1)).getSeqno()).isEqualTo(2L);
        assertThat(requests.get(2)).isInstanceOf(ExecuteBatchDmlRequest.class);
        assertThat(((ExecuteBatchDmlRequest) requests.get(2)).getSeqno()).isEqualTo(3L);
        assertThat(requests.get(3)).isInstanceOf(ExecuteSqlRequest.class);
        assertThat(((ExecuteSqlRequest) requests.get(3)).getSeqno()).isEqualTo(4L);
    }
}

Example 4 with CallbackResponse

use of com.google.cloud.spanner.AsyncResultSet.CallbackResponse in project java-spanner by googleapis.

the class ITAsyncExamplesTest method readAsync.

@Test
public void readAsync() throws Exception {
    final SettableApiFuture<List<String>> future = SettableApiFuture.create();
    try (AsyncResultSet rs = client.singleUse().readAsync(TABLE_NAME, KeySet.all(), ALL_COLUMNS)) {
        rs.setCallback(executor, new ReadyCallback() {

            final List<String> values = new LinkedList<>();

            @Override
            public CallbackResponse cursorReady(AsyncResultSet resultSet) {
                try {
                    while (true) {
                        switch(resultSet.tryNext()) {
                            case DONE:
                                future.set(values);
                                return CallbackResponse.DONE;
                            case NOT_READY:
                                return CallbackResponse.CONTINUE;
                            case OK:
                                values.add(resultSet.getString("StringValue"));
                                break;
                        }
                    }
                } catch (Throwable t) {
                    future.setException(t);
                    return CallbackResponse.DONE;
                }
            }
        });
    }
    assertThat(future.get()).containsExactlyElementsIn(ALL_VALUES_IN_PK_ORDER);
}

Example 5 with CallbackResponse

use of com.google.cloud.spanner.AsyncResultSet.CallbackResponse in project grpc-gcp-java by GoogleCloudPlatform.

the class SpannerClient method singleReadAsync.

public ApiFuture<Void> singleReadAsync() {
    final AsyncResultSet asyncResultSet = dbClient.singleUse().readAsync(TABLE_NAME, READ_KEY_SET, ALL_COLUMNS);
    final ApiFuture<Void> readFuture = asyncResultSet.setCallback(executor, new ReadyCallback() {

        @Override
        public CallbackResponse cursorReady(AsyncResultSet resultSet) {
            while (true) {
                switch(resultSet.tryNext()) {
                    case OK:
                        // System.out.println(resultSet.getString("StringValue"));
                        break;
                    case NOT_READY:
                        return CallbackResponse.CONTINUE;
                    case DONE:
                        return CallbackResponse.DONE;
                    default:
                        throw new IllegalStateException();
                }
            }
        }
    });
    return readFuture;
}