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());
}
}
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();
}
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();
}
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);
}
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.
}
}
Aggregations