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

Example 1 with ExecuteBatchDmlResponse

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

the class TransactionContextImplTest method batchDml.

@SuppressWarnings("unchecked")
private void batchDml(int status) {
    SessionImpl session = mock(SessionImpl.class);
    when(session.getName()).thenReturn("test");
    SpannerRpc rpc = mock(SpannerRpc.class);
    ExecuteBatchDmlResponse response = ExecuteBatchDmlResponse.newBuilder().setStatus(Status.newBuilder().setCode(status).build()).build();
    Statement statement = Statement.of("UPDATE FOO SET BAR=1");
    when(rpc.executeBatchDml(Mockito.any(ExecuteBatchDmlRequest.class), Mockito.anyMap())).thenReturn(response);
    try (TransactionContextImpl impl = TransactionContextImpl.newBuilder().setSession(session).setRpc(rpc).setTransactionId(ByteString.copyFromUtf8("test")).setOptions(Options.fromTransactionOptions()).build()) {
        impl.batchUpdate(Collections.singletonList(statement));
    }
}

Example 2 with ExecuteBatchDmlResponse

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

use of com.google.spanner.v1.ExecuteBatchDmlResponse 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());
    }
}

Example 4 with ExecuteBatchDmlResponse

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

the class TransactionRunnerImplTest method batchDmlException.

@SuppressWarnings("unchecked")
private long[] batchDmlException(int status) {
    Preconditions.checkArgument(status != Code.OK_VALUE);
    TransactionContextImpl transaction = TransactionContextImpl.newBuilder().setSession(session).setTransactionId(ByteString.copyFromUtf8(UUID.randomUUID().toString())).setOptions(Options.fromTransactionOptions()).setRpc(rpc).build();
    when(session.newTransaction(Options.fromTransactionOptions())).thenReturn(transaction);
    when(session.beginTransactionAsync()).thenReturn(ApiFutures.immediateFuture(ByteString.copyFromUtf8(UUID.randomUUID().toString())));
    when(session.getName()).thenReturn(SessionId.of("p", "i", "d", "test").getName());
    TransactionRunnerImpl runner = new TransactionRunnerImpl(session);
    runner.setSpan(mock(Span.class));
    ExecuteBatchDmlResponse response1 = ExecuteBatchDmlResponse.newBuilder().addResultSets(ResultSet.newBuilder().setStats(ResultSetStats.newBuilder().setRowCountExact(1L)).build()).setStatus(com.google.rpc.Status.newBuilder().setCode(status).build()).build();
    ExecuteBatchDmlResponse response2 = ExecuteBatchDmlResponse.newBuilder().addResultSets(ResultSet.newBuilder().setStats(ResultSetStats.newBuilder().setRowCountExact(1L)).build()).addResultSets(ResultSet.newBuilder().setStats(ResultSetStats.newBuilder().setRowCountExact(1L)).build()).setStatus(com.google.rpc.Status.newBuilder().setCode(Code.OK_VALUE).build()).build();
    when(rpc.executeBatchDml(Mockito.any(ExecuteBatchDmlRequest.class), Mockito.anyMap())).thenReturn(response1, response2);
    CommitResponse commitResponse = CommitResponse.newBuilder().setCommitTimestamp(Timestamp.getDefaultInstance()).build();
    when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap())).thenReturn(ApiFutures.immediateFuture(commitResponse));
    final Statement statement = Statement.of("UPDATE FOO SET BAR=1");
    final AtomicInteger numCalls = new AtomicInteger(0);
    long[] updateCount = runner.run(transaction1 -> {
        numCalls.incrementAndGet();
        return transaction1.batchUpdate(Arrays.asList(statement, statement));
    });
    if (status == Code.ABORTED_VALUE) {
        // Assert that the method ran twice because the first response aborted.
        assertThat(numCalls.get()).isEqualTo(2);
    }
    return updateCount;
}

Example 5 with ExecuteBatchDmlResponse

use of com.google.spanner.v1.ExecuteBatchDmlResponse in project grpc-gcp-java by GoogleCloudPlatform.

the class SpannerIntegrationTest method testExecuteBatchDmlFuture.

@Test
public void testExecuteBatchDmlFuture() throws Exception {
    SpannerFutureStub stub = getSpannerFutureStub();
    List<String> futureNames = createFutureSessions(stub);
    for (String futureName : futureNames) {
        TransactionOptions options = TransactionOptions.newBuilder().setReadWrite(TransactionOptions.ReadWrite.getDefaultInstance()).build();
        TransactionSelector selector = TransactionSelector.newBuilder().setBegin(options).build();
        // Will use only one session for the whole batch.
        ListenableFuture<ExecuteBatchDmlResponse> responseFuture = stub.executeBatchDml(ExecuteBatchDmlRequest.newBuilder().setSession(futureName).setTransaction(selector).addStatements(Statement.newBuilder().setSql("select * FROM Users").build()).build());
        // The ChannelRef which is bound with the current affinity key.
        GcpManagedChannel.ChannelRef currentChannel = gcpChannel.affinityKeyToChannelRef.get(futureName);
        // Verify the channel is in use.
        assertEquals(1, currentChannel.getActiveStreamsCount());
        ExecuteBatchDmlResponse response = responseFuture.get();
        assertEquals(0, currentChannel.getActiveStreamsCount());
    }
    deleteFutureSessions(stub, futureNames);
}