Make integration tests work

googleapis · Apr 13, 2018 · 067a190 · 067a190
1 parent f208590
commit 067a190
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 15 deletions.
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
@@ -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
@@ -165,7 +165,7 @@ class SpannerImpl extends BaseService<SpannerOptions> implements Spanner {
   SpannerImpl(SpannerOptions options) {
     this(
         options.getSpannerRpcV1(),
-        GapicSpannerRpc.create(options),
+        options.getGapicSpannerRpc(),
         options.getPrefetchChunks(),
         options);
   }
@@ -1063,7 +1063,9 @@ CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken
                   null,
                   session.options));
 
-              // let resume fail for now
+              // Let resume fail for now. Gapic has its own resume, but in order not 
+              // to introduce too much change at a time, we decide to plumb up
+              // ServerStream first and then figure out how to make resume work
             }
           };
       return new GrpcResultSet(stream, this, queryMode);
@@ -1170,7 +1172,9 @@ CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken
                 null,
                 session.options));
 
-              // let resume fail for now
+              // Let resume fail for now. Gapic has its own resume, but in order not 
+              // to introduce too much change at a time, we decide to plumb up
+              // ServerStream first and then figure out how to make resume work
             }
           };
       GrpcResultSet resultSet =
@@ -2288,17 +2292,32 @@ public CloseableServerStreamIterator(ServerStream<T> 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
@@ -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/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java
@@ -402,8 +402,7 @@ public ServerStream<PartialResultSet> read(
   @Override
   public ServerStream<PartialResultSet> executeQuery(
       ExecuteSqlRequest request, ResultStreamConsumer consumer, @Nullable Map<Option, ?> options) {
-    GrpcCallContext context = GrpcCallContext.createDefault()
-        .withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue());
+    GrpcCallContext context = newCallContext(options, request.getSession());
     return stub.executeStreamingSqlCallable().call(request, context);
   }
 

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
@@ -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<Map<SpannerRpc.Option, Object>> 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<String, String>) anyMap(), optionsCaptor.capture()))
+    when(gapicRpc.createSession(eq(DB_NAME), (Map<String, String>) 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.<BeginTransactionRequest>any(), optionsCaptor.capture()))
+    when(spannerOptions.getGapicSpannerRpc()).thenReturn(gapicRpc);
+    when(gapicRpc.beginTransaction(Mockito.<BeginTransactionRequest>any(), optionsCaptor.capture()))
         .thenReturn(txnMetadata);
 
     BatchReadOnlyTransaction batchTxn = client.batchReadOnlyTransaction(TimestampBound.strong());