Adding a callable that doesn't post lantency to the client metrics

googleapis · Nov 30, 2021 · 0bae71d · 0bae71d
1 parent d0709e8
commit 0bae71d
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 1 deletion.
diff --git a/...-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java b/...-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java
@@ -144,7 +144,7 @@ private void sendPrimeRequests(ManagedChannel managedChannel) throws IOException
       // Prime all of the table ids in parallel
       for (String tableId : tableIds) {
         ApiFuture<Row> f =
-            stub.readRowCallable()
+            stub.readRowRawCallable()
                 .futureCall(Query.create(tableId).rowKey(PRIMING_ROW_KEY).filter(FILTERS.block()));
 
         primeFutures.put(tableId, f);

diff --git a/...d-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/...d-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
@@ -132,6 +132,7 @@ public class EnhancedBigtableStub implements AutoCloseable {
 
   private final ServerStreamingCallable<Query, Row> readRowsCallable;
   private final UnaryCallable<Query, Row> readRowCallable;
+  private final UnaryCallable<Query, Row> readRowRawCallable;
   private final UnaryCallable<String, List<KeyOffset>> sampleRowKeysCallable;
   private final UnaryCallable<RowMutation, Void> mutateRowCallable;
   private final UnaryCallable<BulkMutation, Void> bulkMutateRowsCallable;
@@ -267,6 +268,7 @@ public EnhancedBigtableStub(EnhancedBigtableStubSettings settings, ClientContext
 
     readRowsCallable = createReadRowsCallable(new DefaultRowAdapter());
     readRowCallable = createReadRowCallable(new DefaultRowAdapter());
+    readRowRawCallable = createReadRowRawCallable(new DefaultRowAdapter());
     sampleRowKeysCallable = createSampleRowKeysCallable();
     mutateRowCallable = createMutateRowCallable();
     bulkMutateRowsCallable = createBulkMutateRowsCallable();
@@ -327,6 +329,26 @@ public <RowT> ServerStreamingCallable<Query, RowT> createReadRowsCallable(
     return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
 
+  /**
+   * Creates a callable chain to handle point ReadRows RPCs. The chain will:
+   *
+   *
+   *
+   * <p>NOTE: the caller is responsible for adding tracing & metrics.</p>
+   */
+  public <RowT> UnaryCallable<Query, RowT> createReadRowRawCallable(RowAdapter<RowT> rowAdapter) {
+    ServerStreamingCallable<ReadRowsRequest, RowT> readRowsCallable =
+        createReadRowsBaseCallable(
+            ServerStreamingCallSettings.<ReadRowsRequest, Row>newBuilder()
+                .setRetryableCodes(settings.readRowSettings().getRetryableCodes())
+                .setRetrySettings(settings.readRowSettings().getRetrySettings())
+                .setIdleTimeout(settings.readRowSettings().getRetrySettings().getTotalTimeout())
+                .build(),
+            rowAdapter).withDefaultCallContext(clientContext.getDefaultCallContext());
+
+    return new ReadRowsUserCallable<>(readRowsCallable, requestContext).first();
+  }
+
   /**
    * Creates a callable chain to handle point ReadRows RPCs. The chain will:
    *
@@ -771,6 +793,10 @@ public UnaryCallable<Query, Row> readRowCallable() {
     return readRowCallable;
   }
 
+  public UnaryCallable<Query, Row> readRowRawCallable() {
+    return readRowRawCallable;
+  }
+
   public UnaryCallable<String, List<KeyOffset>> sampleRowKeysCallable() {
     return sampleRowKeysCallable;
   }

diff --git a/...gtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java b/...gtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java
@@ -31,6 +31,7 @@
 import com.google.api.gax.grpc.GrpcTransportChannel;
 import com.google.api.gax.rpc.FixedTransportChannelProvider;
 import com.google.api.gax.rpc.ServerStreamingCallable;
+import com.google.api.gax.rpc.UnaryCallable;
 import com.google.auth.oauth2.ServiceAccountJwtAccessCredentials;
 import com.google.bigtable.v2.BigtableGrpc;
 import com.google.bigtable.v2.MutateRowsRequest;
@@ -262,6 +263,19 @@ public void testCreateReadRowsRawCallable() throws InterruptedException {
     assertThat(fakeDataService.popLastRequest()).isEqualTo(expectedRequest2);
   }
 
+  @Test
+  public void testCreateReadRowRawCallable() throws InterruptedException {
+    UnaryCallable<Query, Row> callable =
+        enhancedBigtableStub.createReadRowRawCallable(new DefaultRowAdapter());
+
+    Query request = Query.create("table-id").rowKey("row-key");
+    callable.call(request);
+
+    ReadRowsRequest expected =
+        request.toProto(RequestContext.create(PROJECT_ID, INSTANCE_ID, APP_PROFILE_ID));
+    assertThat(fakeDataService.popLastRequest()).isEqualTo(expected);
+  }
+
   @Test
   public void testChannelPrimerConfigured() throws IOException {
     EnhancedBigtableStubSettings settings =