diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java index dbb1d705d32c..e07c935b0f78 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java @@ -23,6 +23,7 @@ import com.google.cloud.spanner.SpannerImpl.MultiUseReadOnlyTransaction; import com.google.cloud.spanner.SpannerImpl.SessionImpl; import com.google.cloud.spanner.spi.v1.SpannerRpc; +import com.google.cloud.spanner.spi.v1.GapicSpannerRpc; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; import com.google.protobuf.Struct; @@ -68,7 +69,7 @@ private static class BatchReadOnlyTransactionImpl extends MultiUseReadOnlyTransa super( checkNotNull(session), checkNotNull(bound), - checkNotNull(spanner).getOptions().getSpannerRpcV1(), + checkNotNull(spanner).getOptions().getGapicSpannerRpc(), spanner.getOptions().getPrefetchChunks()); this.sessionName = session.getName(); this.options = session.getOptions(); @@ -81,7 +82,7 @@ private static class BatchReadOnlyTransactionImpl extends MultiUseReadOnlyTransa checkNotNull(session), checkNotNull(batchTransactionId).getTransactionId(), batchTransactionId.getTimestamp(), - checkNotNull(spanner).getOptions().getSpannerRpcV1(), + checkNotNull(spanner).getOptions().getGapicSpannerRpc(), spanner.getOptions().getPrefetchChunks()); this.sessionName = session.getName(); this.options = session.getOptions(); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java index 363b2b5e9d0d..c326572fe346 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java @@ -165,7 +165,7 @@ class SpannerImpl extends BaseService implements Spanner { SpannerImpl(SpannerOptions options) { this( options.getSpannerRpcV1(), - GapicSpannerRpc.create(options), + options.getGapicSpannerRpc(), options.getPrefetchChunks(), options); } @@ -2288,17 +2288,32 @@ public CloseableServerStreamIterator(ServerStream stream) { @Override public boolean hasNext() { - return iterator.hasNext(); + try { + return iterator.hasNext(); + } + catch (Exception e) { + throw SpannerExceptionFactory.newSpannerException(e); + } } @Override public T next() { - return iterator.next(); + try { + return iterator.next(); + } + catch (Exception e) { + throw SpannerExceptionFactory.newSpannerException(e); + } } @Override public void close(@Nullable String message) { - stream.cancel(); + try { + stream.cancel(); + } + catch (Exception e) { + throw SpannerExceptionFactory.newSpannerException(e); + } } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java index b5ca59c1b9c1..a8990112c769 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java @@ -21,6 +21,7 @@ import com.google.cloud.ServiceOptions; import com.google.cloud.ServiceRpc; import com.google.cloud.TransportOptions; +import com.google.cloud.spanner.spi.v1.GapicSpannerRpc; import com.google.cloud.spanner.spi.v1.GrpcSpannerRpc; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.cloud.spanner.spi.SpannerRpcFactory; @@ -343,6 +344,10 @@ protected SpannerRpc getSpannerRpcV1() { return (SpannerRpc) getRpc(); } + protected SpannerRpc getGapicSpannerRpc() { + return GapicSpannerRpc.create(this); + } + @SuppressWarnings("unchecked") @Override public Builder toBuilder() { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BatchClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BatchClientImplTest.java index da9a39cb057d..a17537bfead0 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BatchClientImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BatchClientImplTest.java @@ -48,7 +48,8 @@ public final class BatchClientImplTest { private static final ByteString TXN_ID = ByteString.copyFromUtf8("my-txn"); private static final String TIMESTAMP = "2017-11-15T10:54:20Z"; - @Mock private SpannerRpc rpc; + @Mock private SpannerRpc rawGrpcRpc; + @Mock private SpannerRpc gapicRpc; @Mock private SpannerOptions spannerOptions; @Captor private ArgumentCaptor> optionsCaptor; @Mock private BatchTransactionId txnID; @@ -59,20 +60,20 @@ public final class BatchClientImplTest { public void setUp() { initMocks(this); DatabaseId db = DatabaseId.of(DB_NAME); - SpannerImpl spanner = new SpannerImpl(rpc, rpc, 1, spannerOptions); + SpannerImpl spanner = new SpannerImpl(rawGrpcRpc, gapicRpc, 1, spannerOptions); client = new BatchClientImpl(db, spanner); } @Test public void testBatchReadOnlyTxnWithBound() throws Exception { Session sessionProto = Session.newBuilder().setName(SESSION_NAME).build(); - when(rpc.createSession(eq(DB_NAME), (Map) anyMap(), optionsCaptor.capture())) + when(gapicRpc.createSession(eq(DB_NAME), (Map) anyMap(), optionsCaptor.capture())) .thenReturn(sessionProto); com.google.protobuf.Timestamp timestamp = Timestamps.parse(TIMESTAMP); Transaction txnMetadata = Transaction.newBuilder().setId(TXN_ID).setReadTimestamp(timestamp).build(); - when(spannerOptions.getSpannerRpcV1()).thenReturn(rpc); - when(rpc.beginTransaction(Mockito.any(), optionsCaptor.capture())) + when(spannerOptions.getGapicSpannerRpc()).thenReturn(gapicRpc); + when(gapicRpc.beginTransaction(Mockito.any(), optionsCaptor.capture())) .thenReturn(txnMetadata); BatchReadOnlyTransaction batchTxn = client.batchReadOnlyTransaction(TimestampBound.strong());