Search in sources :

Example 26 with SessionName

use of com.google.spanner.v1.SessionName 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());
    }
}
Also used : TransactionCallable(com.google.cloud.spanner.TransactionRunner.TransactionCallable) Arrays(java.util.Arrays) ClosedException(com.google.cloud.spanner.SpannerImpl.ClosedException) QueryAnalyzeMode(com.google.cloud.spanner.ReadContext.QueryAnalyzeMode) CommitResponse(com.google.spanner.v1.CommitResponse) Timestamp(com.google.cloud.Timestamp) Clock(com.google.cloud.spanner.SessionPool.Clock) PooledSession(com.google.cloud.spanner.SessionPool.PooledSession) Mockito.doThrow(org.mockito.Mockito.doThrow) Empty(com.google.protobuf.Empty) ResultSetStats(com.google.spanner.v1.ResultSetStats) Future(java.util.concurrent.Future) Mockito.doAnswer(org.mockito.Mockito.doAnswer) Map(java.util.Map) Assert.fail(org.junit.Assert.fail) Parameterized(org.junit.runners.Parameterized) ResultStreamConsumer(com.google.cloud.spanner.spi.v1.SpannerRpc.ResultStreamConsumer) PrintWriter(java.io.PrintWriter) PointWithFunction(com.google.cloud.spanner.MetricRegistryTestUtils.PointWithFunction) ApiFutures(com.google.api.core.ApiFutures) Collection(java.util.Collection) NUM_READ_SESSIONS(com.google.cloud.spanner.MetricRegistryConstants.NUM_READ_SESSIONS) Executors(java.util.concurrent.Executors) ByteString(com.google.protobuf.ByteString) CountDownLatch(java.util.concurrent.CountDownLatch) NUM_SESSIONS_BEING_PREPARED(com.google.cloud.spanner.MetricRegistryConstants.NUM_SESSIONS_BEING_PREPARED) List(java.util.List) LabelValue(io.opencensus.metrics.LabelValue) Mockito.atMost(org.mockito.Mockito.atMost) ExecuteSqlRequest(com.google.spanner.v1.ExecuteSqlRequest) Mockito.any(org.mockito.Mockito.any) MetricRegistry(io.opencensus.metrics.MetricRegistry) SPANNER_LABEL_KEYS_WITH_TYPE(com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS_WITH_TYPE) Mockito.eq(org.mockito.Mockito.eq) Mockito.mock(org.mockito.Mockito.mock) MockitoAnnotations.initMocks(org.mockito.MockitoAnnotations.initMocks) ListenableFuture(com.google.common.util.concurrent.ListenableFuture) Mock(org.mockito.Mock) SPANNER_LABEL_KEYS(com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS) Assert.assertThrows(org.junit.Assert.assertThrows) RunWith(org.junit.runner.RunWith) Parameters(org.junit.runners.Parameterized.Parameters) SpannerRpc(com.google.cloud.spanner.spi.v1.SpannerRpc) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) ArrayList(java.util.ArrayList) CommitRequest(com.google.spanner.v1.CommitRequest) PooledSessionFuture(com.google.cloud.spanner.SessionPool.PooledSessionFuture) MetricsRecord(com.google.cloud.spanner.MetricRegistryTestUtils.MetricsRecord) SessionConsumerImpl(com.google.cloud.spanner.SessionPool.SessionConsumerImpl) LinkedList(java.util.LinkedList) ExecutorService(java.util.concurrent.ExecutorService) Before(org.junit.Before) FakeMetricRegistry(com.google.cloud.spanner.MetricRegistryTestUtils.FakeMetricRegistry) Uninterruptibles(com.google.common.util.concurrent.Uninterruptibles) NUM_IN_USE_SESSIONS(com.google.cloud.spanner.MetricRegistryConstants.NUM_IN_USE_SESSIONS) Assert.assertNotNull(org.junit.Assert.assertNotNull) Parameter(org.junit.runners.Parameterized.Parameter) StringWriter(java.io.StringWriter) Assert.assertTrue(org.junit.Assert.assertTrue) Mockito.times(org.mockito.Mockito.times) Test(org.junit.Test) Mockito.when(org.mockito.Mockito.when) Truth.assertThat(com.google.common.truth.Truth.assertThat) Span(io.opencensus.trace.Span) Mockito.verify(org.mockito.Mockito.verify) RollbackRequest(com.google.spanner.v1.RollbackRequest) NUM_WRITE_SESSIONS(com.google.cloud.spanner.MetricRegistryConstants.NUM_WRITE_SESSIONS) TimeUnit(java.util.concurrent.TimeUnit) Mockito(org.mockito.Mockito) Mockito.never(org.mockito.Mockito.never) Mockito.anyInt(org.mockito.Mockito.anyInt) TransactionContextImpl(com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl) ExecuteBatchDmlRequest(com.google.spanner.v1.ExecuteBatchDmlRequest) Collections(java.util.Collections) Assert.assertEquals(org.junit.Assert.assertEquals) SessionConsumer(com.google.cloud.spanner.SessionClient.SessionConsumer) RollbackRequest(com.google.spanner.v1.RollbackRequest) SpannerRpc(com.google.cloud.spanner.spi.v1.SpannerRpc) Span(io.opencensus.trace.Span) SessionConsumer(com.google.cloud.spanner.SessionClient.SessionConsumer) ClosedException(com.google.cloud.spanner.SpannerImpl.ClosedException) CommitRequest(com.google.spanner.v1.CommitRequest) ClosedException(com.google.cloud.spanner.SpannerImpl.ClosedException) TransactionContextImpl(com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl) SessionConsumerImpl(com.google.cloud.spanner.SessionPool.SessionConsumerImpl) ExecuteSqlRequest(com.google.spanner.v1.ExecuteSqlRequest) ResultStreamConsumer(com.google.cloud.spanner.spi.v1.SpannerRpc.ResultStreamConsumer) PooledSessionFuture(com.google.cloud.spanner.SessionPool.PooledSessionFuture) Map(java.util.Map) ExecuteBatchDmlRequest(com.google.spanner.v1.ExecuteBatchDmlRequest) Test(org.junit.Test)

Example 27 with SessionName

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

the class SessionPoolTest method testSessionNotFoundWrite.

@Test
public void testSessionNotFoundWrite() {
    SpannerException sessionNotFound = SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName);
    List<Mutation> mutations = Collections.singletonList(Mutation.newInsertBuilder("FOO").build());
    final SessionImpl closedSession = mockSession();
    when(closedSession.writeWithOptions(mutations)).thenThrow(sessionNotFound);
    final SessionImpl openSession = mockSession();
    com.google.cloud.spanner.CommitResponse response = mock(com.google.cloud.spanner.CommitResponse.class);
    when(response.getCommitTimestamp()).thenReturn(Timestamp.now());
    when(openSession.writeWithOptions(mutations)).thenReturn(response);
    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));
    FakeClock clock = new FakeClock();
    clock.currentTimeMillis = System.currentTimeMillis();
    pool = createPool(clock);
    DatabaseClientImpl impl = new DatabaseClientImpl(pool);
    assertThat(impl.write(mutations)).isNotNull();
}
Also used : TransactionCallable(com.google.cloud.spanner.TransactionRunner.TransactionCallable) Arrays(java.util.Arrays) ClosedException(com.google.cloud.spanner.SpannerImpl.ClosedException) QueryAnalyzeMode(com.google.cloud.spanner.ReadContext.QueryAnalyzeMode) CommitResponse(com.google.spanner.v1.CommitResponse) Timestamp(com.google.cloud.Timestamp) Clock(com.google.cloud.spanner.SessionPool.Clock) PooledSession(com.google.cloud.spanner.SessionPool.PooledSession) Mockito.doThrow(org.mockito.Mockito.doThrow) Empty(com.google.protobuf.Empty) ResultSetStats(com.google.spanner.v1.ResultSetStats) Future(java.util.concurrent.Future) Mockito.doAnswer(org.mockito.Mockito.doAnswer) Map(java.util.Map) Assert.fail(org.junit.Assert.fail) Parameterized(org.junit.runners.Parameterized) ResultStreamConsumer(com.google.cloud.spanner.spi.v1.SpannerRpc.ResultStreamConsumer) PrintWriter(java.io.PrintWriter) PointWithFunction(com.google.cloud.spanner.MetricRegistryTestUtils.PointWithFunction) ApiFutures(com.google.api.core.ApiFutures) Collection(java.util.Collection) NUM_READ_SESSIONS(com.google.cloud.spanner.MetricRegistryConstants.NUM_READ_SESSIONS) Executors(java.util.concurrent.Executors) ByteString(com.google.protobuf.ByteString) CountDownLatch(java.util.concurrent.CountDownLatch) NUM_SESSIONS_BEING_PREPARED(com.google.cloud.spanner.MetricRegistryConstants.NUM_SESSIONS_BEING_PREPARED) List(java.util.List) LabelValue(io.opencensus.metrics.LabelValue) Mockito.atMost(org.mockito.Mockito.atMost) ExecuteSqlRequest(com.google.spanner.v1.ExecuteSqlRequest) Mockito.any(org.mockito.Mockito.any) MetricRegistry(io.opencensus.metrics.MetricRegistry) SPANNER_LABEL_KEYS_WITH_TYPE(com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS_WITH_TYPE) Mockito.eq(org.mockito.Mockito.eq) Mockito.mock(org.mockito.Mockito.mock) MockitoAnnotations.initMocks(org.mockito.MockitoAnnotations.initMocks) ListenableFuture(com.google.common.util.concurrent.ListenableFuture) Mock(org.mockito.Mock) SPANNER_LABEL_KEYS(com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS) Assert.assertThrows(org.junit.Assert.assertThrows) RunWith(org.junit.runner.RunWith) Parameters(org.junit.runners.Parameterized.Parameters) SpannerRpc(com.google.cloud.spanner.spi.v1.SpannerRpc) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) ArrayList(java.util.ArrayList) CommitRequest(com.google.spanner.v1.CommitRequest) PooledSessionFuture(com.google.cloud.spanner.SessionPool.PooledSessionFuture) MetricsRecord(com.google.cloud.spanner.MetricRegistryTestUtils.MetricsRecord) SessionConsumerImpl(com.google.cloud.spanner.SessionPool.SessionConsumerImpl) LinkedList(java.util.LinkedList) ExecutorService(java.util.concurrent.ExecutorService) Before(org.junit.Before) FakeMetricRegistry(com.google.cloud.spanner.MetricRegistryTestUtils.FakeMetricRegistry) Uninterruptibles(com.google.common.util.concurrent.Uninterruptibles) NUM_IN_USE_SESSIONS(com.google.cloud.spanner.MetricRegistryConstants.NUM_IN_USE_SESSIONS) Assert.assertNotNull(org.junit.Assert.assertNotNull) Parameter(org.junit.runners.Parameterized.Parameter) StringWriter(java.io.StringWriter) Assert.assertTrue(org.junit.Assert.assertTrue) Mockito.times(org.mockito.Mockito.times) Test(org.junit.Test) Mockito.when(org.mockito.Mockito.when) Truth.assertThat(com.google.common.truth.Truth.assertThat) Span(io.opencensus.trace.Span) Mockito.verify(org.mockito.Mockito.verify) RollbackRequest(com.google.spanner.v1.RollbackRequest) NUM_WRITE_SESSIONS(com.google.cloud.spanner.MetricRegistryConstants.NUM_WRITE_SESSIONS) TimeUnit(java.util.concurrent.TimeUnit) Mockito(org.mockito.Mockito) Mockito.never(org.mockito.Mockito.never) Mockito.anyInt(org.mockito.Mockito.anyInt) TransactionContextImpl(com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl) ExecuteBatchDmlRequest(com.google.spanner.v1.ExecuteBatchDmlRequest) Collections(java.util.Collections) Assert.assertEquals(org.junit.Assert.assertEquals) SessionConsumer(com.google.cloud.spanner.SessionClient.SessionConsumer) SessionConsumer(com.google.cloud.spanner.SessionClient.SessionConsumer) SessionConsumerImpl(com.google.cloud.spanner.SessionPool.SessionConsumerImpl) Test(org.junit.Test)

Example 28 with SessionName

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

the class SessionImplTest method setUp.

@SuppressWarnings("unchecked")
@Before
public void setUp() {
    MockitoAnnotations.initMocks(this);
    when(spannerOptions.getNumChannels()).thenReturn(4);
    when(spannerOptions.getPrefetchChunks()).thenReturn(1);
    when(spannerOptions.getRetrySettings()).thenReturn(RetrySettings.newBuilder().build());
    when(spannerOptions.getClock()).thenReturn(NanoClock.getDefaultClock());
    when(spannerOptions.getSessionLabels()).thenReturn(Collections.emptyMap());
    GrpcTransportOptions transportOptions = mock(GrpcTransportOptions.class);
    when(transportOptions.getExecutorFactory()).thenReturn(mock(ExecutorFactory.class));
    when(spannerOptions.getTransportOptions()).thenReturn(transportOptions);
    when(spannerOptions.getSessionPoolOptions()).thenReturn(mock(SessionPoolOptions.class));
    @SuppressWarnings("resource") SpannerImpl spanner = new SpannerImpl(rpc, spannerOptions);
    String dbName = "projects/p1/instances/i1/databases/d1";
    String sessionName = dbName + "/sessions/s1";
    DatabaseId db = DatabaseId.of(dbName);
    Session sessionProto = Session.newBuilder().setName(sessionName).build();
    Mockito.when(rpc.createSession(Mockito.eq(dbName), Mockito.anyMap(), optionsCaptor.capture())).thenReturn(sessionProto);
    Transaction txn = Transaction.newBuilder().setId(ByteString.copyFromUtf8("TEST")).build();
    Mockito.when(rpc.beginTransactionAsync(Mockito.any(BeginTransactionRequest.class), Mockito.any(Map.class))).thenReturn(ApiFutures.immediateFuture(txn));
    CommitResponse commitResponse = CommitResponse.newBuilder().setCommitTimestamp(com.google.protobuf.Timestamp.getDefaultInstance()).build();
    Mockito.when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.any(Map.class))).thenReturn(ApiFutures.immediateFuture(commitResponse));
    Mockito.when(rpc.rollbackAsync(Mockito.any(RollbackRequest.class), Mockito.anyMap())).thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance()));
    session = spanner.getSessionClient(db).createSession();
    ((SessionImpl) session).setCurrentSpan(mock(Span.class));
    // We expect the same options, "options", on all calls on "session".
    options = optionsCaptor.getValue();
}
Also used : CommitRequest(com.google.spanner.v1.CommitRequest) CommitResponse(com.google.spanner.v1.CommitResponse) ByteString(com.google.protobuf.ByteString) RollbackRequest(com.google.spanner.v1.RollbackRequest) Span(io.opencensus.trace.Span) GrpcTransportOptions(com.google.cloud.grpc.GrpcTransportOptions) Transaction(com.google.spanner.v1.Transaction) ExecutorFactory(com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory) BeginTransactionRequest(com.google.spanner.v1.BeginTransactionRequest) Map(java.util.Map) Session(com.google.spanner.v1.Session) Before(org.junit.Before)

Example 29 with SessionName

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

the class SessionPoolTest method testSessionNotFoundPartitionedUpdate.

@Test
public void testSessionNotFoundPartitionedUpdate() {
    SpannerException sessionNotFound = SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName);
    Statement statement = Statement.of("UPDATE FOO SET BAR=1 WHERE 1=1");
    final SessionImpl closedSession = mockSession();
    when(closedSession.executePartitionedUpdate(statement)).thenThrow(sessionNotFound);
    final SessionImpl openSession = mockSession();
    when(openSession.executePartitionedUpdate(statement)).thenReturn(1L);
    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));
    FakeClock clock = new FakeClock();
    clock.currentTimeMillis = System.currentTimeMillis();
    pool = createPool(clock);
    DatabaseClientImpl impl = new DatabaseClientImpl(pool);
    assertThat(impl.executePartitionedUpdate(statement)).isEqualTo(1L);
}
Also used : TransactionCallable(com.google.cloud.spanner.TransactionRunner.TransactionCallable) Arrays(java.util.Arrays) ClosedException(com.google.cloud.spanner.SpannerImpl.ClosedException) QueryAnalyzeMode(com.google.cloud.spanner.ReadContext.QueryAnalyzeMode) CommitResponse(com.google.spanner.v1.CommitResponse) Timestamp(com.google.cloud.Timestamp) Clock(com.google.cloud.spanner.SessionPool.Clock) PooledSession(com.google.cloud.spanner.SessionPool.PooledSession) Mockito.doThrow(org.mockito.Mockito.doThrow) Empty(com.google.protobuf.Empty) ResultSetStats(com.google.spanner.v1.ResultSetStats) Future(java.util.concurrent.Future) Mockito.doAnswer(org.mockito.Mockito.doAnswer) Map(java.util.Map) Assert.fail(org.junit.Assert.fail) Parameterized(org.junit.runners.Parameterized) ResultStreamConsumer(com.google.cloud.spanner.spi.v1.SpannerRpc.ResultStreamConsumer) PrintWriter(java.io.PrintWriter) PointWithFunction(com.google.cloud.spanner.MetricRegistryTestUtils.PointWithFunction) ApiFutures(com.google.api.core.ApiFutures) Collection(java.util.Collection) NUM_READ_SESSIONS(com.google.cloud.spanner.MetricRegistryConstants.NUM_READ_SESSIONS) Executors(java.util.concurrent.Executors) ByteString(com.google.protobuf.ByteString) CountDownLatch(java.util.concurrent.CountDownLatch) NUM_SESSIONS_BEING_PREPARED(com.google.cloud.spanner.MetricRegistryConstants.NUM_SESSIONS_BEING_PREPARED) List(java.util.List) LabelValue(io.opencensus.metrics.LabelValue) Mockito.atMost(org.mockito.Mockito.atMost) ExecuteSqlRequest(com.google.spanner.v1.ExecuteSqlRequest) Mockito.any(org.mockito.Mockito.any) MetricRegistry(io.opencensus.metrics.MetricRegistry) SPANNER_LABEL_KEYS_WITH_TYPE(com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS_WITH_TYPE) Mockito.eq(org.mockito.Mockito.eq) Mockito.mock(org.mockito.Mockito.mock) MockitoAnnotations.initMocks(org.mockito.MockitoAnnotations.initMocks) ListenableFuture(com.google.common.util.concurrent.ListenableFuture) Mock(org.mockito.Mock) SPANNER_LABEL_KEYS(com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS) Assert.assertThrows(org.junit.Assert.assertThrows) RunWith(org.junit.runner.RunWith) Parameters(org.junit.runners.Parameterized.Parameters) SpannerRpc(com.google.cloud.spanner.spi.v1.SpannerRpc) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) ArrayList(java.util.ArrayList) CommitRequest(com.google.spanner.v1.CommitRequest) PooledSessionFuture(com.google.cloud.spanner.SessionPool.PooledSessionFuture) MetricsRecord(com.google.cloud.spanner.MetricRegistryTestUtils.MetricsRecord) SessionConsumerImpl(com.google.cloud.spanner.SessionPool.SessionConsumerImpl) LinkedList(java.util.LinkedList) ExecutorService(java.util.concurrent.ExecutorService) Before(org.junit.Before) FakeMetricRegistry(com.google.cloud.spanner.MetricRegistryTestUtils.FakeMetricRegistry) Uninterruptibles(com.google.common.util.concurrent.Uninterruptibles) NUM_IN_USE_SESSIONS(com.google.cloud.spanner.MetricRegistryConstants.NUM_IN_USE_SESSIONS) Assert.assertNotNull(org.junit.Assert.assertNotNull) Parameter(org.junit.runners.Parameterized.Parameter) StringWriter(java.io.StringWriter) Assert.assertTrue(org.junit.Assert.assertTrue) Mockito.times(org.mockito.Mockito.times) Test(org.junit.Test) Mockito.when(org.mockito.Mockito.when) Truth.assertThat(com.google.common.truth.Truth.assertThat) Span(io.opencensus.trace.Span) Mockito.verify(org.mockito.Mockito.verify) RollbackRequest(com.google.spanner.v1.RollbackRequest) NUM_WRITE_SESSIONS(com.google.cloud.spanner.MetricRegistryConstants.NUM_WRITE_SESSIONS) TimeUnit(java.util.concurrent.TimeUnit) Mockito(org.mockito.Mockito) Mockito.never(org.mockito.Mockito.never) Mockito.anyInt(org.mockito.Mockito.anyInt) TransactionContextImpl(com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl) ExecuteBatchDmlRequest(com.google.spanner.v1.ExecuteBatchDmlRequest) Collections(java.util.Collections) Assert.assertEquals(org.junit.Assert.assertEquals) SessionConsumer(com.google.cloud.spanner.SessionClient.SessionConsumer) SessionConsumer(com.google.cloud.spanner.SessionClient.SessionConsumer) SessionConsumerImpl(com.google.cloud.spanner.SessionPool.SessionConsumerImpl) Test(org.junit.Test)

Example 30 with SessionName

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

the class SpannerClientTest method commitExceptionTest2.

@Test
public void commitExceptionTest2() throws Exception {
    StatusRuntimeException exception = new StatusRuntimeException(io.grpc.Status.INVALID_ARGUMENT);
    mockSpanner.addException(exception);
    try {
        SessionName session = SessionName.of("[PROJECT]", "[INSTANCE]", "[DATABASE]", "[SESSION]");
        TransactionOptions singleUseTransaction = TransactionOptions.newBuilder().build();
        List<Mutation> mutations = new ArrayList<>();
        client.commit(session, singleUseTransaction, mutations);
        Assert.fail("No exception raised");
    } catch (InvalidArgumentException e) {
    // Expected exception.
    }
}
Also used : InvalidArgumentException(com.google.api.gax.rpc.InvalidArgumentException) TransactionOptions(com.google.spanner.v1.TransactionOptions) StatusRuntimeException(io.grpc.StatusRuntimeException) ArrayList(java.util.ArrayList) Mutation(com.google.spanner.v1.Mutation) SessionName(com.google.spanner.v1.SessionName) Test(org.junit.Test)

Aggregations

Test (org.junit.Test)18 ByteString (com.google.protobuf.ByteString)13 SessionName (com.google.spanner.v1.SessionName)12 ArrayList (java.util.ArrayList)10 CommitRequest (com.google.spanner.v1.CommitRequest)9 CommitResponse (com.google.spanner.v1.CommitResponse)9 RollbackRequest (com.google.spanner.v1.RollbackRequest)9 Empty (com.google.protobuf.Empty)8 Before (org.junit.Before)8 ApiFutures (com.google.api.core.ApiFutures)6 Timestamp (com.google.cloud.Timestamp)6 NUM_IN_USE_SESSIONS (com.google.cloud.spanner.MetricRegistryConstants.NUM_IN_USE_SESSIONS)6 NUM_READ_SESSIONS (com.google.cloud.spanner.MetricRegistryConstants.NUM_READ_SESSIONS)6 NUM_SESSIONS_BEING_PREPARED (com.google.cloud.spanner.MetricRegistryConstants.NUM_SESSIONS_BEING_PREPARED)6 NUM_WRITE_SESSIONS (com.google.cloud.spanner.MetricRegistryConstants.NUM_WRITE_SESSIONS)6 SPANNER_LABEL_KEYS (com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS)6 SPANNER_LABEL_KEYS_WITH_TYPE (com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS_WITH_TYPE)6 FakeMetricRegistry (com.google.cloud.spanner.MetricRegistryTestUtils.FakeMetricRegistry)6 MetricsRecord (com.google.cloud.spanner.MetricRegistryTestUtils.MetricsRecord)6 PointWithFunction (com.google.cloud.spanner.MetricRegistryTestUtils.PointWithFunction)6