Examples with SpannerRpc com.google.cloud.spanner.spi.v1.SpannerRpc used on opensource projects

Example 1 with SpannerRpc

use of com.google.cloud.spanner.spi.v1.SpannerRpc 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 SpannerRpc

use of com.google.cloud.spanner.spi.v1.SpannerRpc in project java-spanner by googleapis.

the class SessionPoolTest method testSessionNotFoundReadWriteTransaction.

@SuppressWarnings("unchecked")
@Test
public void testSessionNotFoundReadWriteTransaction() {
    final Statement queryStatement = Statement.of("SELECT 1");
    final Statement updateStatement = Statement.of("UPDATE FOO SET BAR=1 WHERE ID=2");
    final SpannerException sessionNotFound = SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName);
    for (ReadWriteTransactionTestStatementType statementType : ReadWriteTransactionTestStatementType.values()) {
        final ReadWriteTransactionTestStatementType executeStatementType = statementType;
        SpannerRpc.StreamingCall closedStreamingCall = mock(SpannerRpc.StreamingCall.class);
        doThrow(sessionNotFound).when(closedStreamingCall).request(Mockito.anyInt());
        SpannerRpc rpc = mock(SpannerRpc.class);
        when(rpc.asyncDeleteSession(Mockito.anyString(), Mockito.anyMap())).thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance()));
        when(rpc.executeQuery(any(ExecuteSqlRequest.class), any(ResultStreamConsumer.class), any(Map.class))).thenReturn(closedStreamingCall);
        when(rpc.executeQuery(any(ExecuteSqlRequest.class), any(Map.class))).thenThrow(sessionNotFound);
        when(rpc.executeBatchDml(any(ExecuteBatchDmlRequest.class), any(Map.class))).thenThrow(sessionNotFound);
        when(rpc.commitAsync(any(CommitRequest.class), any(Map.class))).thenReturn(ApiFutures.<CommitResponse>immediateFailedFuture(sessionNotFound));
        when(rpc.rollbackAsync(any(RollbackRequest.class), any(Map.class))).thenReturn(ApiFutures.<Empty>immediateFailedFuture(sessionNotFound));
        final SessionImpl closedSession = mock(SessionImpl.class);
        when(closedSession.getName()).thenReturn("projects/dummy/instances/dummy/database/dummy/sessions/session-closed");
        final TransactionContextImpl closedTransactionContext = TransactionContextImpl.newBuilder().setSession(closedSession).setOptions(Options.fromTransactionOptions()).setRpc(rpc).build();
        when(closedSession.asyncClose()).thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance()));
        when(closedSession.newTransaction(Options.fromTransactionOptions())).thenReturn(closedTransactionContext);
        when(closedSession.beginTransactionAsync()).thenThrow(sessionNotFound);
        TransactionRunnerImpl closedTransactionRunner = new TransactionRunnerImpl(closedSession);
        closedTransactionRunner.setSpan(mock(Span.class));
        when(closedSession.readWriteTransaction()).thenReturn(closedTransactionRunner);
        final SessionImpl openSession = mock(SessionImpl.class);
        when(openSession.asyncClose()).thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance()));
        when(openSession.getName()).thenReturn("projects/dummy/instances/dummy/database/dummy/sessions/session-open");
        final TransactionContextImpl openTransactionContext = mock(TransactionContextImpl.class);
        when(openSession.newTransaction(Options.fromTransactionOptions())).thenReturn(openTransactionContext);
        when(openSession.beginTransactionAsync()).thenReturn(ApiFutures.immediateFuture(ByteString.copyFromUtf8("open-txn")));
        TransactionRunnerImpl openTransactionRunner = new TransactionRunnerImpl(openSession);
        openTransactionRunner.setSpan(mock(Span.class));
        when(openSession.readWriteTransaction()).thenReturn(openTransactionRunner);
        ResultSet openResultSet = mock(ResultSet.class);
        when(openResultSet.next()).thenReturn(true, false);
        ResultSet planResultSet = mock(ResultSet.class);
        when(planResultSet.getStats()).thenReturn(ResultSetStats.getDefaultInstance());
        when(openTransactionContext.executeQuery(queryStatement)).thenReturn(openResultSet);
        when(openTransactionContext.analyzeQuery(queryStatement, QueryAnalyzeMode.PLAN)).thenReturn(planResultSet);
        when(openTransactionContext.executeUpdate(updateStatement)).thenReturn(1L);
        when(openTransactionContext.batchUpdate(Arrays.asList(updateStatement, updateStatement))).thenReturn(new long[] { 1L, 1L });
        SpannerImpl spanner = mock(SpannerImpl.class);
        SessionClient sessionClient = mock(SessionClient.class);
        when(spanner.getSessionClient(db)).thenReturn(sessionClient);
        doAnswer(invocation -> {
            executor.submit(() -> {
                SessionConsumerImpl consumer = invocation.getArgument(2, SessionConsumerImpl.class);
                consumer.onSessionReady(closedSession);
            });
            return null;
        }).doAnswer(invocation -> {
            executor.submit(() -> {
                SessionConsumerImpl consumer = invocation.getArgument(2, SessionConsumerImpl.class);
                consumer.onSessionReady(openSession);
            });
            return null;
        }).when(sessionClient).asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class));
        SessionPoolOptions options = SessionPoolOptions.newBuilder().setMinSessions(// The pool should not auto-create any sessions
        0).setMaxSessions(2).setIncStep(1).setBlockIfPoolExhausted().build();
        SpannerOptions spannerOptions = mock(SpannerOptions.class);
        when(spannerOptions.getSessionPoolOptions()).thenReturn(options);
        when(spannerOptions.getNumChannels()).thenReturn(4);
        when(spanner.getOptions()).thenReturn(spannerOptions);
        SessionPool pool = SessionPool.createPool(options, new TestExecutorFactory(), spanner.getSessionClient(db));
        try (PooledSessionFuture readWriteSession = pool.getSession()) {
            TransactionRunner runner = readWriteSession.readWriteTransaction();
            try {
                runner.run(new TransactionCallable<Integer>() {

                    private int callNumber = 0;

                    @Override
                    public Integer run(TransactionContext transaction) {
                        callNumber++;
                        if (callNumber == 1) {
                            assertThat(transaction).isEqualTo(closedTransactionContext);
                        } else {
                            assertThat(transaction).isEqualTo(openTransactionContext);
                        }
                        switch(executeStatementType) {
                            case QUERY:
                                ResultSet resultSet = transaction.executeQuery(queryStatement);
                                assertThat(resultSet.next()).isTrue();
                                break;
                            case ANALYZE:
                                ResultSet planResultSet = transaction.analyzeQuery(queryStatement, QueryAnalyzeMode.PLAN);
                                assertThat(planResultSet.next()).isFalse();
                                assertThat(planResultSet.getStats()).isNotNull();
                                break;
                            case UPDATE:
                                long updateCount = transaction.executeUpdate(updateStatement);
                                assertThat(updateCount).isEqualTo(1L);
                                break;
                            case BATCH_UPDATE:
                                long[] updateCounts = transaction.batchUpdate(Arrays.asList(updateStatement, updateStatement));
                                assertThat(updateCounts).isEqualTo(new long[] { 1L, 1L });
                                break;
                            case WRITE:
                                transaction.buffer(Mutation.delete("FOO", Key.of(1L)));
                                break;
                            case EXCEPTION:
                                throw new RuntimeException("rollback at call " + callNumber);
                            default:
                                fail("Unknown statement type: " + executeStatementType);
                        }
                        return callNumber;
                    }
                });
            } catch (Exception e) {
                // The rollback will also cause a SessionNotFoundException, but this is caught, logged
                // and further ignored by the library, meaning that the session will not be re-created
                // for retry. Hence rollback at call 1.
                assertThat(executeStatementType).isEqualTo(ReadWriteTransactionTestStatementType.EXCEPTION);
                assertThat(e.getMessage()).contains("rollback at call 1");
            }
        }
        pool.closeAsync(new SpannerImpl.ClosedException());
    }
}

Example 3 with SpannerRpc

use of com.google.cloud.spanner.spi.v1.SpannerRpc in project java-spanner by googleapis.

the class TransactionManagerImplTest method usesPreparedTransaction.

@SuppressWarnings("unchecked")
@Test
public void usesPreparedTransaction() {
    SpannerOptions options = mock(SpannerOptions.class);
    when(options.getNumChannels()).thenReturn(4);
    GrpcTransportOptions transportOptions = mock(GrpcTransportOptions.class);
    when(transportOptions.getExecutorFactory()).thenReturn(new TestExecutorFactory());
    when(options.getTransportOptions()).thenReturn(transportOptions);
    SessionPoolOptions sessionPoolOptions = SessionPoolOptions.newBuilder().setMinSessions(0).setIncStep(1).build();
    when(options.getSessionPoolOptions()).thenReturn(sessionPoolOptions);
    when(options.getSessionLabels()).thenReturn(Collections.emptyMap());
    SpannerRpc rpc = mock(SpannerRpc.class);
    when(rpc.asyncDeleteSession(Mockito.anyString(), Mockito.anyMap())).thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance()));
    when(rpc.batchCreateSessions(Mockito.anyString(), Mockito.eq(1), Mockito.anyMap(), Mockito.anyMap())).thenAnswer(invocation -> Collections.singletonList(Session.newBuilder().setName(invocation.getArguments()[0] + "/sessions/1").setCreateTime(com.google.protobuf.Timestamp.newBuilder().setSeconds(System.currentTimeMillis() * 1000)).build()));
    when(rpc.beginTransactionAsync(Mockito.any(BeginTransactionRequest.class), Mockito.anyMap())).thenAnswer(invocation -> ApiFutures.immediateFuture(Transaction.newBuilder().setId(ByteString.copyFromUtf8(UUID.randomUUID().toString())).build()));
    when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap())).thenAnswer(invocation -> ApiFutures.immediateFuture(com.google.spanner.v1.CommitResponse.newBuilder().setCommitTimestamp(com.google.protobuf.Timestamp.newBuilder().setSeconds(System.currentTimeMillis() * 1000)).build()));
    DatabaseId db = DatabaseId.of("test", "test", "test");
    try (SpannerImpl spanner = new SpannerImpl(rpc, options)) {
        DatabaseClient client = spanner.getDatabaseClient(db);
        try (TransactionManager mgr = client.transactionManager()) {
            mgr.begin();
            mgr.commit();
        }
        verify(rpc, times(1)).beginTransactionAsync(Mockito.any(BeginTransactionRequest.class), Mockito.anyMap());
    }
}

Example 4 with SpannerRpc

use of com.google.cloud.spanner.spi.v1.SpannerRpc in project java-spanner by googleapis.

the class TransactionManagerImplTest method inlineBegin.

@SuppressWarnings({ "unchecked", "resource" })
@Test
public void inlineBegin() {
    SpannerOptions options = mock(SpannerOptions.class);
    when(options.getNumChannels()).thenReturn(4);
    GrpcTransportOptions transportOptions = mock(GrpcTransportOptions.class);
    when(transportOptions.getExecutorFactory()).thenReturn(new TestExecutorFactory());
    when(options.getTransportOptions()).thenReturn(transportOptions);
    SessionPoolOptions sessionPoolOptions = SessionPoolOptions.newBuilder().setMinSessions(0).setIncStep(1).build();
    when(options.getSessionPoolOptions()).thenReturn(sessionPoolOptions);
    when(options.getSessionLabels()).thenReturn(Collections.emptyMap());
    when(options.getDefaultQueryOptions(Mockito.any(DatabaseId.class))).thenReturn(QueryOptions.getDefaultInstance());
    SpannerRpc rpc = mock(SpannerRpc.class);
    when(rpc.asyncDeleteSession(Mockito.anyString(), Mockito.anyMap())).thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance()));
    when(rpc.batchCreateSessions(Mockito.anyString(), Mockito.eq(1), Mockito.anyMap(), Mockito.anyMap())).thenAnswer(invocation -> Collections.singletonList(Session.newBuilder().setName(invocation.getArguments()[0] + "/sessions/1").setCreateTime(com.google.protobuf.Timestamp.newBuilder().setSeconds(System.currentTimeMillis() * 1000)).build()));
    when(rpc.beginTransactionAsync(Mockito.any(BeginTransactionRequest.class), Mockito.anyMap())).thenAnswer(invocation -> ApiFutures.immediateFuture(Transaction.newBuilder().setId(ByteString.copyFromUtf8(UUID.randomUUID().toString())).build()));
    final AtomicInteger transactionsStarted = new AtomicInteger();
    when(rpc.executeQuery(Mockito.any(ExecuteSqlRequest.class), Mockito.anyMap())).thenAnswer(invocation -> {
        ResultSet.Builder builder = ResultSet.newBuilder().setStats(ResultSetStats.newBuilder().setRowCountExact(1L).build());
        ExecuteSqlRequest request = invocation.getArgument(0, ExecuteSqlRequest.class);
        if (request.getTransaction() != null && request.getTransaction().hasBegin()) {
            transactionsStarted.incrementAndGet();
            builder.setMetadata(ResultSetMetadata.newBuilder().setTransaction(Transaction.newBuilder().setId(ByteString.copyFromUtf8("test-tx")).build()).build());
        }
        return builder.build();
    });
    when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap())).thenAnswer(invocation -> ApiFutures.immediateFuture(com.google.spanner.v1.CommitResponse.newBuilder().setCommitTimestamp(com.google.protobuf.Timestamp.newBuilder().setSeconds(System.currentTimeMillis() * 1000)).build()));
    DatabaseId db = DatabaseId.of("test", "test", "test");
    try (SpannerImpl spanner = new SpannerImpl(rpc, options)) {
        DatabaseClient client = spanner.getDatabaseClient(db);
        try (TransactionManager mgr = client.transactionManager()) {
            TransactionContext tx = mgr.begin();
            while (true) {
                try {
                    tx.executeUpdate(Statement.of("UPDATE FOO SET BAR=1"));
                    tx.executeUpdate(Statement.of("UPDATE FOO SET BAZ=2"));
                    mgr.commit();
                    break;
                } catch (AbortedException e) {
                    tx = mgr.resetForRetry();
                }
            }
        }
        // BeginTransaction should not be called, as we are inlining it with the ExecuteSql request.
        verify(rpc, Mockito.never()).beginTransaction(Mockito.any(BeginTransactionRequest.class), Mockito.anyMap());
        // We should have 2 ExecuteSql requests.
        verify(rpc, times(2)).executeQuery(Mockito.any(ExecuteSqlRequest.class), Mockito.anyMap());
        // But only 1 with a BeginTransaction.
        assertThat(transactionsStarted.get()).isEqualTo(1);
    }
}

Example 5 with SpannerRpc

use of com.google.cloud.spanner.spi.v1.SpannerRpc in project java-spanner by googleapis.

the class TransactionRunnerImplTest method usesPreparedTransaction.

@SuppressWarnings("unchecked")
@Test
public void usesPreparedTransaction() {
    SpannerOptions options = mock(SpannerOptions.class);
    when(options.getNumChannels()).thenReturn(4);
    GrpcTransportOptions transportOptions = mock(GrpcTransportOptions.class);
    when(transportOptions.getExecutorFactory()).thenReturn(new TestExecutorFactory());
    when(options.getTransportOptions()).thenReturn(transportOptions);
    SessionPoolOptions sessionPoolOptions = SessionPoolOptions.newBuilder().setMinSessions(0).setIncStep(1).build();
    when(options.getSessionPoolOptions()).thenReturn(sessionPoolOptions);
    when(options.getSessionLabels()).thenReturn(Collections.emptyMap());
    SpannerRpc rpc = mock(SpannerRpc.class);
    when(rpc.asyncDeleteSession(Mockito.anyString(), Mockito.anyMap())).thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance()));
    when(rpc.batchCreateSessions(Mockito.anyString(), Mockito.eq(1), Mockito.anyMap(), Mockito.anyMap())).thenAnswer(invocation -> Collections.singletonList(Session.newBuilder().setName(invocation.getArguments()[0] + "/sessions/1").setCreateTime(Timestamp.newBuilder().setSeconds(System.currentTimeMillis() * 1000)).build()));
    when(rpc.beginTransactionAsync(Mockito.any(BeginTransactionRequest.class), Mockito.anyMap())).thenAnswer(invocation -> ApiFutures.immediateFuture(Transaction.newBuilder().setId(ByteString.copyFromUtf8(UUID.randomUUID().toString())).build()));
    when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap())).thenAnswer(invocation -> ApiFutures.immediateFuture(CommitResponse.newBuilder().setCommitTimestamp(Timestamp.newBuilder().setSeconds(System.currentTimeMillis() * 1000)).build()));
    DatabaseId db = DatabaseId.of("test", "test", "test");
    try (SpannerImpl spanner = new SpannerImpl(rpc, options)) {
        DatabaseClient client = spanner.getDatabaseClient(db);
        client.readWriteTransaction().run(transaction -> null);
        verify(rpc, times(1)).beginTransactionAsync(Mockito.any(BeginTransactionRequest.class), Mockito.anyMap());
    }
}