Examples with SpannerApiFutures.get com.google.cloud.spanner.SpannerApiFutures.get used on opensource projects

Example 1 with SpannerApiFutures.get

use of com.google.cloud.spanner.SpannerApiFutures.get in project java-spanner by googleapis.

the class ConnectionAsyncApiTest method testExecuteQueryAsync.

private void testExecuteQueryAsync(Function<Connection, Void> connectionConfigurator, boolean executeAsStatement) {
    ApiFuture<Void> res;
    try (Connection connection = createConnection()) {
        connectionConfigurator.apply(connection);
        for (boolean timeout : new boolean[] { true, false }) {
            final AtomicInteger rowCount = new AtomicInteger();
            final AtomicBoolean receivedTimeout = new AtomicBoolean();
            if (timeout) {
                mockSpanner.setExecuteStreamingSqlExecutionTime(SimulatedExecutionTime.ofMinimumAndRandomTime(10, 0));
                connection.setStatementTimeout(1L, TimeUnit.NANOSECONDS);
            } else {
                mockSpanner.removeAllExecutionTimes();
                connection.clearStatementTimeout();
            }
            try (AsyncResultSet rs = executeAsStatement ? connection.executeAsync(SELECT_RANDOM_STATEMENT).getResultSetAsync() : connection.executeQueryAsync(SELECT_RANDOM_STATEMENT)) {
                res = rs.setCallback(executor, resultSet -> {
                    try {
                        while (true) {
                            switch(resultSet.tryNext()) {
                                case OK:
                                    rowCount.incrementAndGet();
                                    break;
                                case DONE:
                                    return CallbackResponse.DONE;
                                case NOT_READY:
                                    return CallbackResponse.CONTINUE;
                            }
                        }
                    } catch (SpannerException e) {
                        receivedTimeout.set(e.getErrorCode() == ErrorCode.DEADLINE_EXCEEDED);
                        throw e;
                    }
                });
            }
            try {
                SpannerApiFutures.get(res);
                assertThat(rowCount.get()).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT);
                if (connection.isReadOnly() || !connection.isInTransaction()) {
                    assertThat(connection.getReadTimestamp()).isNotNull();
                }
                assertThat(timeout).isFalse();
            } catch (SpannerException e) {
                assertThat(e.getSuppressed()).hasLength(1);
                assertThat(e.getSuppressed()[0].getMessage()).contains(SELECT_RANDOM_STATEMENT.getSql());
                assertThat(e.getErrorCode()).isEqualTo(ErrorCode.DEADLINE_EXCEEDED);
                assertThat(timeout).isTrue();
                assertThat(receivedTimeout.get()).isTrue();
                // invalidate that transaction.
                if (!connection.isReadOnly() && connection.isInTransaction()) {
                    connection.clearStatementTimeout();
                    connection.rollback();
                }
            }
        }
    }
}

Example 2 with SpannerApiFutures.get

use of com.google.cloud.spanner.SpannerApiFutures.get in project java-spanner by googleapis.

the class ConnectionAsyncApiTest method testExecuteQueryAsyncIsNonBlocking.

private void testExecuteQueryAsyncIsNonBlocking(Function<Connection, Void> connectionConfigurator) {
    ApiFuture<Void> res;
    final AtomicInteger rowCount = new AtomicInteger();
    mockSpanner.freeze();
    try (Connection connection = createConnection()) {
        connectionConfigurator.apply(connection);
        try (AsyncResultSet rs = connection.executeQueryAsync(SELECT_RANDOM_STATEMENT)) {
            res = rs.setCallback(executor, resultSet -> {
                while (true) {
                    switch(resultSet.tryNext()) {
                        case OK:
                            rowCount.incrementAndGet();
                            break;
                        case DONE:
                            return CallbackResponse.DONE;
                        case NOT_READY:
                            return CallbackResponse.CONTINUE;
                    }
                }
            });
            mockSpanner.unfreeze();
        }
        SpannerApiFutures.get(res);
        assertThat(rowCount.get()).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT);
    }
}