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

Example 1 with ReadyCallback

use of com.google.cloud.spanner.AsyncResultSet.ReadyCallback 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 ReadyCallback

use of com.google.cloud.spanner.AsyncResultSet.ReadyCallback 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 ReadyCallback

use of com.google.cloud.spanner.AsyncResultSet.ReadyCallback 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 ReadyCallback

use of com.google.cloud.spanner.AsyncResultSet.ReadyCallback 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 ReadyCallback

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

the class AsyncQueryExample method asyncQuery.

// Execute a query asynchronously and process the results in a callback.
static void asyncQuery(DatabaseClient client) throws InterruptedException, ExecutionException, TimeoutException {
    ApiFuture<Void> finished;
    ExecutorService executor = Executors.newSingleThreadExecutor();
    try (AsyncResultSet resultSet = client.singleUse().executeQueryAsync(Statement.of("SELECT SingerId, AlbumId, AlbumTitle FROM Albums"))) {
        // Setting a callback will automatically start the iteration over the results of the query
        // using the specified executor. The callback will be called at least once. The returned
        // ApiFuture is done when the callback has returned DONE and all resources used by the
        // AsyncResultSet have been released.
        finished = resultSet.setCallback(executor, new ReadyCallback() {

            @Override
            public CallbackResponse cursorReady(AsyncResultSet resultSet) {
                try {
                    while (true) {
                        switch(resultSet.tryNext()) {
                            // OK: There is a row ready.
                            case OK:
                                System.out.printf("%d %d %s%n", resultSet.getLong(0), resultSet.getLong(1), resultSet.getString(2));
                                break;
                            // DONE: There are no more rows in the result set.
                            case DONE:
                                return CallbackResponse.DONE;
                            // NOT_READY: There are currently no more rows in the buffer.
                            case NOT_READY:
                                return CallbackResponse.CONTINUE;
                            default:
                                throw new IllegalStateException();
                        }
                    }
                } catch (SpannerException e) {
                    System.out.printf("Error in callback: %s%n", e.getMessage());
                    return CallbackResponse.DONE;
                }
            }
        });
    }
    // This ApiFuture is done when the callback has returned DONE and all resources of the
    // asynchronous result set have been released.
    finished.get(30L, TimeUnit.SECONDS);
    executor.shutdown();
}