Examples with ExecuteBatchDmlRequest com.google.spanner.v1.ExecuteBatchDmlRequest used on opensource projects

Example 1 with ExecuteBatchDmlRequest

use of com.google.spanner.v1.ExecuteBatchDmlRequest in project java-spanner by googleapis.

the class DatabaseClientImplTest method testBatchUpdateWithPriority.

@Test
public void testBatchUpdateWithPriority() {
    DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
    TransactionRunner runner = client.readWriteTransaction();
    runner.run(transaction -> transaction.batchUpdate(Collections.singletonList(UPDATE_STATEMENT), Options.priority(RpcPriority.HIGH)));
    List<ExecuteBatchDmlRequest> requests = mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class);
    assertThat(requests).hasSize(1);
    ExecuteBatchDmlRequest request = requests.get(0);
    assertNotNull(request.getRequestOptions());
    assertEquals(Priority.PRIORITY_HIGH, request.getRequestOptions().getPriority());
}

Example 2 with ExecuteBatchDmlRequest

use of com.google.spanner.v1.ExecuteBatchDmlRequest 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 3 with ExecuteBatchDmlRequest

use of com.google.spanner.v1.ExecuteBatchDmlRequest in project java-spanner by googleapis.

the class SpannerClientTest method executeBatchDmlTest.

@Test
public void executeBatchDmlTest() throws Exception {
    ExecuteBatchDmlResponse expectedResponse = ExecuteBatchDmlResponse.newBuilder().addAllResultSets(new ArrayList<ResultSet>()).setStatus(Status.newBuilder().build()).build();
    mockSpanner.addResponse(expectedResponse);
    ExecuteBatchDmlRequest request = ExecuteBatchDmlRequest.newBuilder().setSession(SessionName.of("[PROJECT]", "[INSTANCE]", "[DATABASE]", "[SESSION]").toString()).setTransaction(TransactionSelector.newBuilder().build()).addAllStatements(new ArrayList<ExecuteBatchDmlRequest.Statement>()).setSeqno(109325920).setRequestOptions(RequestOptions.newBuilder().build()).build();
    ExecuteBatchDmlResponse actualResponse = client.executeBatchDml(request);
    Assert.assertEquals(expectedResponse, actualResponse);
    List<AbstractMessage> actualRequests = mockSpanner.getRequests();
    Assert.assertEquals(1, actualRequests.size());
    ExecuteBatchDmlRequest actualRequest = ((ExecuteBatchDmlRequest) actualRequests.get(0));
    Assert.assertEquals(request.getSession(), actualRequest.getSession());
    Assert.assertEquals(request.getTransaction(), actualRequest.getTransaction());
    Assert.assertEquals(request.getStatementsList(), actualRequest.getStatementsList());
    Assert.assertEquals(request.getSeqno(), actualRequest.getSeqno());
    Assert.assertEquals(request.getRequestOptions(), actualRequest.getRequestOptions());
    Assert.assertTrue(channelProvider.isHeaderSent(ApiClientHeaderProvider.getDefaultApiClientHeaderKey(), GaxGrpcProperties.getDefaultApiClientHeaderPattern()));
}

Example 4 with ExecuteBatchDmlRequest

use of com.google.spanner.v1.ExecuteBatchDmlRequest in project java-spanner by googleapis.

the class DatabaseClientImplTest method testBatchUpdateWithTag.

@Test
public void testBatchUpdateWithTag() {
    DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
    TransactionRunner runner = client.readWriteTransaction(Options.tag("app=spanner,env=test,action=txn"));
    runner.run(transaction -> transaction.batchUpdate(Collections.singletonList(UPDATE_STATEMENT), Options.tag("app=spanner,env=test,action=batch")));
    List<ExecuteBatchDmlRequest> requests = mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class);
    assertThat(requests).hasSize(1);
    ExecuteBatchDmlRequest request = requests.get(0);
    assertNotNull(request.getRequestOptions());
    assertThat(request.getRequestOptions().getRequestTag()).isEqualTo("app=spanner,env=test,action=batch");
    assertThat(request.getRequestOptions().getTransactionTag()).isEqualTo("app=spanner,env=test,action=txn");
}

Example 5 with ExecuteBatchDmlRequest

use of com.google.spanner.v1.ExecuteBatchDmlRequest in project java-spanner by googleapis.

the class MockSpannerServiceImpl method executeBatchDml.

@Override
public void executeBatchDml(ExecuteBatchDmlRequest request, StreamObserver<ExecuteBatchDmlResponse> responseObserver) {
    requests.add(request);
    Preconditions.checkNotNull(request.getSession());
    Session session = sessions.get(request.getSession());
    if (session == null) {
        setSessionNotFound(request.getSession(), responseObserver);
        return;
    }
    sessionLastUsed.put(session.getName(), Instant.now());
    try {
        executeBatchDmlExecutionTime.simulateExecutionTime(exceptions, stickyGlobalExceptions, freezeLock);
        // Get or start transaction
        ByteString transactionId = getTransactionId(session, request.getTransaction());
        if (isPartitionedDmlTransaction(transactionId)) {
            throw Status.FAILED_PRECONDITION.withDescription("This transaction is a partitioned DML transaction and cannot be used for batch DML updates.").asRuntimeException();
        }
        simulateAbort(session, transactionId);
        List<StatementResult> results = new ArrayList<>();
        com.google.rpc.Status status = com.google.rpc.Status.newBuilder().setCode(Code.OK_VALUE).build();
        resultLoop: for (com.google.spanner.v1.ExecuteBatchDmlRequest.Statement statement : request.getStatementsList()) {
            try {
                Statement spannerStatement = buildStatement(statement.getSql(), statement.getParamTypesMap(), statement.getParams());
                StatementResult res = getResult(spannerStatement);
                switch(res.getType()) {
                    case EXCEPTION:
                        status = com.google.rpc.Status.newBuilder().setCode(res.getException().getStatus().getCode().value()).setMessage(res.getException().getMessage()).build();
                        break resultLoop;
                    case RESULT_SET:
                        throw Status.INVALID_ARGUMENT.withDescription("Not a DML statement: " + statement.getSql()).asRuntimeException();
                    case UPDATE_COUNT:
                        results.add(res);
                        break;
                    default:
                        throw new IllegalStateException("Unknown result type: " + res.getType());
                }
            } catch (StatusRuntimeException e) {
                status = com.google.rpc.Status.newBuilder().setCode(e.getStatus().getCode().value()).setMessage(e.getMessage()).build();
                break;
            } catch (Exception e) {
                status = com.google.rpc.Status.newBuilder().setCode(Code.UNKNOWN_VALUE).setMessage(e.getMessage()).build();
                break;
            }
        }
        ExecuteBatchDmlResponse.Builder builder = ExecuteBatchDmlResponse.newBuilder();
        for (StatementResult res : results) {
            builder.addResultSets(ResultSet.newBuilder().setStats(ResultSetStats.newBuilder().setRowCountExact(res.getUpdateCount()).build()).setMetadata(ResultSetMetadata.newBuilder().setTransaction(ignoreNextInlineBeginRequest.getAndSet(false) ? Transaction.getDefaultInstance() : Transaction.newBuilder().setId(transactionId).build()).build()).build());
        }
        builder.setStatus(status);
        responseObserver.onNext(builder.build());
        responseObserver.onCompleted();
    } catch (StatusRuntimeException e) {
        responseObserver.onError(e);
    } catch (Throwable t) {
        responseObserver.onError(Status.INTERNAL.asRuntimeException());
    }
}