fix: modify code base on feedback

googleapis · May 22, 2020 · 90c66bb · 90c66bb
1 parent e6d5188
commit 90c66bb
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 27 deletions.
diff --git a/google-cloud-bigquerystorage/pom.xml b/google-cloud-bigquerystorage/pom.xml
@@ -179,7 +179,6 @@
     <dependency>
       <groupId>com.google.api.grpc</groupId>
       <artifactId>grpc-google-cloud-bigquerystorage-v1</artifactId>
-      <scope>test</scope>
     </dependency>
     <!-- Need testing utility classes for generated gRPC clients tests -->
     <dependency>

diff --git a/...age/src/main/java/com/google/cloud/bigquery/storage/v1/stub/EnhancedBigQueryReadStub.java b/...age/src/main/java/com/google/cloud/bigquery/storage/v1/stub/EnhancedBigQueryReadStub.java
@@ -17,13 +17,20 @@
 
 import com.google.api.core.InternalApi;
 import com.google.api.gax.core.BackgroundResource;
+import com.google.api.gax.grpc.GrpcCallSettings;
+import com.google.api.gax.grpc.GrpcRawCallableFactory;
 import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
 import com.google.api.gax.retrying.ScheduledRetryingExecutor;
 import com.google.api.gax.retrying.StreamingRetryAlgorithm;
+import com.google.api.gax.rpc.Callables;
 import com.google.api.gax.rpc.ClientContext;
+import com.google.api.gax.rpc.RequestParamsExtractor;
 import com.google.api.gax.rpc.ServerStreamingCallSettings;
 import com.google.api.gax.rpc.ServerStreamingCallable;
 import com.google.api.gax.rpc.UnaryCallable;
+import com.google.api.gax.tracing.SpanName;
+import com.google.api.gax.tracing.TracedServerStreamingCallable;
+import com.google.cloud.bigquery.storage.v1.BigQueryReadGrpc;
 import com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest;
 import com.google.cloud.bigquery.storage.v1.ReadRowsRequest;
 import com.google.cloud.bigquery.storage.v1.ReadRowsResponse;
@@ -32,7 +39,9 @@
 import com.google.cloud.bigquery.storage.v1.SplitReadStreamResponse;
 import com.google.cloud.bigquery.storage.v1.stub.readrows.ApiResultRetryAlgorithm;
 import com.google.cloud.bigquery.storage.v1.stub.readrows.ReadRowsRetryingCallable;
+import com.google.common.collect.ImmutableMap;
 import java.io.IOException;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -41,6 +50,8 @@
  * <p>This class is for advanced usage and reflects the underlying API directly.
  */
 public class EnhancedBigQueryReadStub implements BackgroundResource {
+
+  private static final String TRACING_OUTER_CLIENT_NAME = "BigQueryStorage";
   private final GrpcBigQueryReadStub stub;
   private final BigQueryReadStubSettings stubSettings;
   private final ClientContext context;
@@ -94,9 +105,22 @@ public UnaryCallable<CreateReadSessionRequest, ReadSession> createReadSessionCal
 
   public ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> readRowsCallable() {
     ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> innerCallable =
-        stub.readRowsCallable();
+        GrpcRawCallableFactory.createServerStreamingCallable(
+            GrpcCallSettings.<ReadRowsRequest, ReadRowsResponse>newBuilder()
+                .setMethodDescriptor(BigQueryReadGrpc.getReadRowsMethod())
+                .setParamsExtractor(
+                    new RequestParamsExtractor<ReadRowsRequest>() {
+                      @Override
+                      public Map<String, String> extract(ReadRowsRequest request) {
+                        return ImmutableMap.of(
+                            "read_stream", String.valueOf(request.getReadStream()));
+                      }
+                    })
+                .build(),
+            stubSettings.readRowsSettings().getRetryableCodes());
     ServerStreamingCallSettings<ReadRowsRequest, ReadRowsResponse> callSettings =
         stubSettings.readRowsSettings();
+
     StreamingRetryAlgorithm<Void> retryAlgorithm =
         new StreamingRetryAlgorithm<>(
             new ApiResultRetryAlgorithm<Void>(),
@@ -105,11 +129,23 @@ public ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> readRowsCallab
     ScheduledRetryingExecutor<Void> retryingExecutor =
         new ScheduledRetryingExecutor<>(retryAlgorithm, context.getExecutor());
 
-    return new ReadRowsRetryingCallable(
-        context.getDefaultCallContext(),
-        innerCallable,
-        retryingExecutor,
-        callSettings.getResumptionStrategy());
+    if (context.getStreamWatchdog() != null) {
+      innerCallable = Callables.watched(innerCallable, callSettings, context);
+    }
+
+    ReadRowsRetryingCallable outerCallable =
+        new ReadRowsRetryingCallable(
+            context.getDefaultCallContext(),
+            innerCallable,
+            retryingExecutor,
+            callSettings.getResumptionStrategy());
+
+    ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> traced =
+        new TracedServerStreamingCallable<>(
+            outerCallable,
+            context.getTracerFactory(),
+            SpanName.of(TRACING_OUTER_CLIENT_NAME, "ReadRows"));
+    return traced.withDefaultCallContext(context.getDefaultCallContext());
   }
 
   public UnaryCallable<SplitReadStreamRequest, SplitReadStreamResponse> splitReadStreamCallable() {

diff --git a/...main/java/com/google/cloud/bigquery/storage/v1/stub/readrows/ApiResultRetryAlgorithm.java b/...main/java/com/google/cloud/bigquery/storage/v1/stub/readrows/ApiResultRetryAlgorithm.java
@@ -32,34 +32,34 @@ public class ApiResultRetryAlgorithm<ResponseT> implements ResultRetryAlgorithm<
   @Override
   public TimedAttemptSettings createNextAttempt(
       Throwable prevThrowable, ResponseT prevResponse, TimedAttemptSettings prevSettings) {
-    if (prevThrowable != null && isRetryable(prevThrowable)) {
-      return TimedAttemptSettings.newBuilder()
-          .setGlobalSettings(prevSettings.getGlobalSettings())
-          .setRetryDelay(prevSettings.getRetryDelay())
-          .setRpcTimeout(prevSettings.getRpcTimeout())
-          .setRandomizedRetryDelay(DEADLINE_SLEEP_DURATION)
-          .setAttemptCount(prevSettings.getAttemptCount() + 1)
-          .setFirstAttemptStartTimeNanos(prevSettings.getFirstAttemptStartTimeNanos())
-          .build();
+    if (prevThrowable != null) {
+      Status status = Status.fromThrowable(prevThrowable);
+      if (status.getCode() == Status.Code.INTERNAL
+          && status.getDescription() != null
+          && status.getDescription().equals("Received unexpected EOS on DATA frame from server")) {
+        return TimedAttemptSettings.newBuilder()
+            .setGlobalSettings(prevSettings.getGlobalSettings())
+            .setRetryDelay(prevSettings.getRetryDelay())
+            .setRpcTimeout(prevSettings.getRpcTimeout())
+            .setRandomizedRetryDelay(DEADLINE_SLEEP_DURATION)
+            .setAttemptCount(prevSettings.getAttemptCount() + 1)
+            .setFirstAttemptStartTimeNanos(prevSettings.getFirstAttemptStartTimeNanos())
+            .build();
+      }
     }
     return null;
   }
 
   @Override
   public boolean shouldRetry(Throwable prevThrowable, ResponseT prevResponse) {
     if (prevThrowable != null) {
-      return isRetryable(prevThrowable);
+      Status status = Status.fromThrowable(prevThrowable);
+      if (status.getCode() == Status.Code.INTERNAL
+          && status.getDescription() != null
+          && status.getDescription().equals("Received unexpected EOS on DATA frame from server")) {
+        return true;
+      }
     }
     return (prevThrowable instanceof ApiException) && ((ApiException) prevThrowable).isRetryable();
   }
-
-  private boolean isRetryable(Throwable prevThrowable) {
-    Status status = Status.fromThrowable(prevThrowable);
-    if (status.getCode() == Status.Code.INTERNAL
-        && status.getDescription() != null
-        && status.getDescription().equals("Received unexpected EOS on DATA frame from server")) {
-      return true;
-    }
-    return false;
-  }
 }