Examples with UPDATE_STATEMENT com.google.cloud.spanner.MockSpannerTestUtil.UPDATE_STATEMENT used on opensource projects

Example 6 with UPDATE_STATEMENT

use of com.google.cloud.spanner.MockSpannerTestUtil.UPDATE_STATEMENT in project java-spanner by googleapis.

the class AsyncTransactionManagerTest method asyncTransactionManagerWithBatchUpdateCommitAborted.

@Test
public void asyncTransactionManagerWithBatchUpdateCommitAborted() throws Exception {
    try (AsyncTransactionManager manager = clientWithEmptySessionPool().transactionManagerAsync()) {
        // Temporarily set the result of the update to 2 rows.
        mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT + 1L));
        final AtomicInteger attempt = new AtomicInteger();
        TransactionContextFuture txn = manager.beginAsync();
        while (true) {
            try {
                AsyncTransactionStep<Void, long[]> updateCounts = txn.then((ignored1, ignored2) -> {
                    if (attempt.get() > 0) {
                        // Set the result of the update statement back to 1 row.
                        mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT));
                    }
                    return ApiFutures.<Void>immediateFuture(null);
                }, executor).then((transactionContext, ignored) -> transactionContext.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)), executor);
                updateCounts.then((transaction, ignored) -> {
                    if (attempt.incrementAndGet() == 1) {
                        mockSpanner.abortTransaction(transaction);
                    }
                    return ApiFutures.immediateFuture(null);
                }, executor).commitAsync().get();
                assertThat(updateCounts.get()).asList().containsExactly(UPDATE_COUNT, UPDATE_COUNT);
                assertThat(attempt.get()).isEqualTo(2);
                break;
            } catch (AbortedException e) {
                txn = manager.resetForRetryAsync();
            }
        }
    } finally {
        mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT));
    }
    assertThat(mockSpanner.getRequestTypes()).containsExactly(BatchCreateSessionsRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class, BeginTransactionRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class);
}

Example 7 with UPDATE_STATEMENT

use of com.google.cloud.spanner.MockSpannerTestUtil.UPDATE_STATEMENT in project java-spanner by googleapis.

the class AsyncTransactionManagerTest method asyncTransactionManagerUpdateAbortedWithoutGettingResult.

@Test
public void asyncTransactionManagerUpdateAbortedWithoutGettingResult() throws Exception {
    final AtomicInteger attempt = new AtomicInteger();
    try (AsyncTransactionManager manager = clientWithEmptySessionPool().transactionManagerAsync()) {
        TransactionContextFuture transactionContextFuture = manager.beginAsync();
        while (true) {
            try {
                CommitTimestampFuture commitTimestampFuture = transactionContextFuture.then((transaction, ignored) -> {
                    if (attempt.incrementAndGet() == 1) {
                        mockSpanner.abortNextStatement();
                    }
                    // This update statement will be aborted, but the error will not
                    // propagated to the transaction runner and cause the transaction to
                    // retry. Instead, the commit call will do that.
                    transaction.executeUpdateAsync(UPDATE_STATEMENT);
                    // has actually finished successfully.
                    return ApiFutures.immediateFuture(null);
                }, executor).commitAsync();
                assertThat(commitTimestampFuture.get()).isNotNull();
                assertThat(attempt.get()).isEqualTo(2);
                // The server may receive 1 or 2 commit requests depending on whether the call to
                // commitAsync() already knows that the transaction has aborted. If it does, it will not
                // attempt to call the Commit RPC and instead directly propagate the Aborted error.
                assertThat(mockSpanner.getRequestTypes()).containsAtLeast(BatchCreateSessionsRequest.class, ExecuteSqlRequest.class, // The retry will use a BeginTransaction RPC.
                BeginTransactionRequest.class, ExecuteSqlRequest.class, CommitRequest.class);
                break;
            } catch (AbortedException e) {
                transactionContextFuture = manager.resetForRetryAsync();
            }
        }
    }
}