remove the setting and make a new callable

Change-Id: I0122fa68ec5de3d20fff7a42fec898e0c3c6597b

Author: mutianf

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java

@@ -151,6 +151,7 @@ public class EnhancedBigtableStub implements AutoCloseable {
   private final ServerStreamingCallable<Query, Row> readRowsCallable;
 
   private final ServerStreamingCallable<Query, Row> skipLargeRowsCallable;
+  private final ServerStreamingCallable<Query, Row> deferLargeRowsCallable;
 
   private final UnaryCallable<Query, Row> readRowCallable;
   private final UnaryCallable<Query, List<Row>> bulkReadRowsCallable;
@@ -172,34 +173,24 @@ public class EnhancedBigtableStub implements AutoCloseable {
   private final ExecuteQueryCallable executeQueryCallable;
   private final UnaryCallable<PrepareQueryRequest, PrepareResponse> prepareQueryCallable;
 
-  private final boolean failOnLargeRows;
-
   public static EnhancedBigtableStub create(EnhancedBigtableStubSettings settings)
       throws IOException {
     BigtableClientContext bigtableClientContext = BigtableClientContext.create(settings);
-    return new EnhancedBigtableStub(
-        settings.getPerOpSettings(), bigtableClientContext, settings.isFailOnLargeRows());
+    return new EnhancedBigtableStub(settings.getPerOpSettings(), bigtableClientContext);
   }
 
   public EnhancedBigtableStub(
       ClientOperationSettings perOpSettings, BigtableClientContext clientContext) {
-    this(perOpSettings, clientContext, false);
-  }
-
-  public EnhancedBigtableStub(
-      ClientOperationSettings perOpSettings,
-      BigtableClientContext clientContext,
-      boolean failOnLargeRows) {
     this.perOpSettings = perOpSettings;
     this.bigtableClientContext = clientContext;
     this.requestContext = RequestContext.create(clientContext.getClientInfo());
     this.bulkMutationFlowController =
         new FlowController(perOpSettings.bulkMutateRowsSettings.getDynamicFlowControlSettings());
     this.bulkMutationDynamicFlowControlStats = new DynamicFlowControlStats();
-    this.failOnLargeRows = failOnLargeRows;
 
     readRowsCallable = createReadRowsCallable(new DefaultRowAdapter());
     skipLargeRowsCallable = createSkipLargeRowsCallable(new DefaultRowAdapter());
+    deferLargeRowsCallable = createDeferLargeRowsCallable(new DefaultRowAdapter());
     readRowCallable = createReadRowCallable(new DefaultRowAdapter());
     bulkReadRowsCallable = createBulkReadRowsCallable(new DefaultRowAdapter());
     sampleRowKeysCallable = createSampleRowKeysCallable();
@@ -420,6 +411,16 @@ private <ReqT, RowT> ServerStreamingCallable<ReadRowsRequest, RowT> createReadRo
    */
   public <ReqT, RowT> ServerStreamingCallable<Query, RowT> createSkipLargeRowsCallable(
       RowAdapter<RowT> rowAdapter) {
+    return createLargeRowsCallable(rowAdapter, false);
+  }
+
+  public <ReqT, RowT> ServerStreamingCallable<Query, RowT> createDeferLargeRowsCallable(
+      RowAdapter<RowT> rowAdapter) {
+    return createLargeRowsCallable(rowAdapter, true);
+  }
+
+  private <ReqT, RowT> ServerStreamingCallable<Query, RowT> createLargeRowsCallable(
+      RowAdapter<RowT> rowAdapter, boolean failOnLargeRows) {
 
     ServerStreamingCallSettings<ReqT, Row> readRowsSettings =
         (ServerStreamingCallSettings<ReqT, Row>) perOpSettings.readRowsSettings;
@@ -1251,6 +1252,14 @@ public ServerStreamingCallable<Query, Row> skipLargeRowsCallable() {
     return skipLargeRowsCallable;
   }
 
+  /**
+   * Returns a streaming read rows callable that defers large rows and throws an exception at the
+   * end
+   */
+  public ServerStreamingCallable<Query, Row> deferLargeRowsCallable() {
+    return deferLargeRowsCallable;
+  }
+
   /** Return a point read callable */
   public UnaryCallable<Query, Row> readRowCallable() {
     return readRowCallable;

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java

@@ -148,7 +148,6 @@ public class EnhancedBigtableStubSettings extends StubSettings<EnhancedBigtableS
   @Nullable private final String metricsEndpoint;
   private final boolean areInternalMetricsEnabled;
   private final String jwtAudience;
-  private final boolean failOnLargeRows;
 
   @InternalApi
   public enum DirectPathConfig {
@@ -173,7 +172,6 @@ private EnhancedBigtableStubSettings(Builder builder) {
     metricsEndpoint = builder.metricsEndpoint;
     areInternalMetricsEnabled = builder.areInternalMetricsEnabled;
     jwtAudience = builder.jwtAudience;
-    failOnLargeRows = builder.failOnLargeRows;
 
     this.sessionsEnabled = builder.sessionsEnabled;
 
@@ -268,10 +266,6 @@ public boolean areInternalMetricsEnabled() {
     return areInternalMetricsEnabled;
   }
 
-  public boolean isFailOnLargeRows() {
-    return failOnLargeRows;
-  }
-
   @InternalApi
   public ClientOperationSettings getPerOpSettings() {
     return perOpSettings;
@@ -617,7 +611,6 @@ public static class Builder extends StubSettings.Builder<EnhancedBigtableStubSet
     private boolean isRefreshingChannel;
     private String jwtAudience;
     private boolean sessionsEnabled = true;
-    private boolean failOnLargeRows;
 
     private final ClientOperationSettings.Builder perOpSettings;
 
@@ -646,7 +639,6 @@ private Builder() {
       this.areInternalMetricsEnabled = true;
       this.jwtAudience = DEFAULT_DATA_JWT_AUDIENCE;
       this.sessionsEnabled = !SESSIONS_DISABLE_ENV_VAR;
-      this.failOnLargeRows = false;
 
       // Defaults provider
       BigtableStubSettings.Builder baseDefaults = BigtableStubSettings.newBuilder();
@@ -683,7 +675,6 @@ private Builder(EnhancedBigtableStubSettings settings) {
       jwtAudience = settings.jwtAudience;
       this.directPathConfig = settings.getDirectPathConfig();
       sessionsEnabled = settings.sessionsEnabled;
-      failOnLargeRows = settings.failOnLargeRows;
 
       this.perOpSettings = new ClientOperationSettings.Builder(settings.perOpSettings);
 
@@ -927,15 +918,6 @@ public boolean getEnableRetryInfo() {
       return true;
     }
 
-    public Builder setFailOnLargeRows(boolean failOnLargeRows) {
-      this.failOnLargeRows = failOnLargeRows;
-      return this;
-    }
-
-    public boolean isFailOnLargeRows() {
-      return failOnLargeRows;
-    }
-
     /** Returns the builder for the settings used for calls to readRows. */
     public ServerStreamingCallSettings.Builder<Query, Row> readRowsSettings() {
       return perOpSettings.readRowsSettings;
@@ -1073,7 +1055,6 @@ public String toString() {
         .add("jwtAudience", jwtAudience)
         .add("directPathConfig", getDirectPathConfig().toString())
         .add("sessionsEnabled", sessionsEnabled)
-        .add("failOnLargeRows", failOnLargeRows)
         .add("parent", super.toString())
         .toString();
   }

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java

@@ -82,7 +82,6 @@ public void settingsAreNotLostTest() {
     WatchdogProvider watchdogProvider = Mockito.mock(WatchdogProvider.class);
     Duration watchdogInterval = Duration.ofSeconds(12);
     String metricsEndpoint = "test-endpoint:443";
-    boolean failOnLargeRows = true;
 
     EnhancedBigtableStubSettings.Builder builder =
         EnhancedBigtableStubSettings.newBuilder()
@@ -94,8 +93,7 @@ public void settingsAreNotLostTest() {
             .setCredentialsProvider(credentialsProvider)
             .setStreamWatchdogProvider(watchdogProvider)
             .setStreamWatchdogCheckInterval(watchdogInterval)
-            .setMetricsEndpoint(metricsEndpoint)
-            .setFailOnLargeRows(failOnLargeRows);
+            .setMetricsEndpoint(metricsEndpoint);
 
     verifyBuilder(
         builder,
@@ -107,8 +105,7 @@ public void settingsAreNotLostTest() {
         credentialsProvider,
         watchdogProvider,
         watchdogInterval,
-        metricsEndpoint,
-        failOnLargeRows);
+        metricsEndpoint);
     verifySettings(
         builder.build(),
         projectId,
@@ -119,8 +116,7 @@ public void settingsAreNotLostTest() {
         credentialsProvider,
         watchdogProvider,
         watchdogInterval,
-        metricsEndpoint,
-        failOnLargeRows);
+        metricsEndpoint);
     verifyBuilder(
         builder.build().toBuilder(),
         projectId,
@@ -131,8 +127,7 @@ public void settingsAreNotLostTest() {
         credentialsProvider,
         watchdogProvider,
         watchdogInterval,
-        metricsEndpoint,
-        failOnLargeRows);
+        metricsEndpoint);
   }
 
   private void verifyBuilder(
@@ -145,8 +140,7 @@ private void verifyBuilder(
       CredentialsProvider credentialsProvider,
       WatchdogProvider watchdogProvider,
       Duration watchdogInterval,
-      String metricsEndpoint,
-      boolean failOnLargeRows) {
+      String metricsEndpoint) {
     assertThat(builder.getProjectId()).isEqualTo(projectId);
     assertThat(builder.getInstanceId()).isEqualTo(instanceId);
     assertThat(builder.getAppProfileId()).isEqualTo(appProfileId);
@@ -156,7 +150,6 @@ private void verifyBuilder(
     assertThat(builder.getStreamWatchdogProvider()).isSameInstanceAs(watchdogProvider);
     assertThat(builder.getStreamWatchdogCheckInterval()).isEqualTo(watchdogInterval);
     assertThat(builder.getMetricsEndpoint()).isEqualTo(metricsEndpoint);
-    assertThat(builder.isFailOnLargeRows()).isEqualTo(failOnLargeRows);
   }
 
   private void verifySettings(
@@ -169,8 +162,7 @@ private void verifySettings(
       CredentialsProvider credentialsProvider,
       WatchdogProvider watchdogProvider,
       Duration watchdogInterval,
-      String metricsEndpoint,
-      boolean failOnLargeRows) {
+      String metricsEndpoint) {
     assertThat(settings.getProjectId()).isEqualTo(projectId);
     assertThat(settings.getInstanceId()).isEqualTo(instanceId);
     assertThat(settings.getAppProfileId()).isEqualTo(appProfileId);
@@ -180,7 +172,6 @@ private void verifySettings(
     assertThat(settings.getStreamWatchdogProvider()).isSameInstanceAs(watchdogProvider);
     assertThat(settings.getStreamWatchdogCheckInterval()).isEqualTo(watchdogInterval);
     assertThat(settings.getMetricsEndpoint()).isEqualTo(metricsEndpoint);
-    assertThat(settings.isFailOnLargeRows()).isEqualTo(failOnLargeRows);
   }
 
   @Test
@@ -916,7 +907,6 @@ public void isRefreshingChannelFalseValueTest() {
     "jwtAudience",
     "directPathConfig",
     "sessionsEnabled",
-    "failOnLargeRows",
   };
 
   @Test

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsRetryTest.java

@@ -39,6 +39,7 @@
 import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
 import com.google.cloud.bigtable.data.v2.models.Row;
 import com.google.cloud.bigtable.data.v2.models.TableId;
+import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub;
 import com.google.cloud.bigtable.data.v2.stub.metrics.NoopMetricsProvider;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
@@ -851,10 +852,9 @@ public void readRowsFailsAtEndIfLargeRowsEncountered() throws Exception {
         .stubSettings()
         .setTransportChannelProvider(
             FixedTransportChannelProvider.create(
-                GrpcTransportChannel.create(serverRule.getChannel())))
-        .setFailOnLargeRows(true);
+                GrpcTransportChannel.create(serverRule.getChannel())));
 
-    BigtableDataClient failureClient = BigtableDataClient.create(settingsBuilder.build());
+    EnhancedBigtableStub stub = EnhancedBigtableStub.create(settingsBuilder.stubSettings().build());
 
     service.expectations.add(
         RpcExpectation.create()
@@ -873,7 +873,7 @@ public void readRowsFailsAtEndIfLargeRowsEncountered() throws Exception {
 
     try {
       ServerStream<Row> actualRows =
-          failureClient.skipLargeRowsCallable().call(Query.create(TABLE_ID).range("r1", "r5"));
+          stub.deferLargeRowsCallable().call(Query.create(TABLE_ID).range("r1", "r5"));
       for (Row row : actualRows) {
         // consume
       }
@@ -886,7 +886,7 @@ public void readRowsFailsAtEndIfLargeRowsEncountered() throws Exception {
           .isEqualTo(com.google.api.gax.rpc.StatusCode.Code.FAILED_PRECONDITION.name());
       assertThat(lre.getLargeRowKeys()).containsExactly(ByteString.copyFromUtf8("r2"));
     } finally {
-      failureClient.close();
+      stub.close();
     }
   }
 
@@ -906,10 +906,9 @@ public void readRowsAppendsLargeRowsToMidStreamError() throws Exception {
         .stubSettings()
         .setTransportChannelProvider(
             FixedTransportChannelProvider.create(
-                GrpcTransportChannel.create(serverRule.getChannel())))
-        .setFailOnLargeRows(true);
+                GrpcTransportChannel.create(serverRule.getChannel())));
 
-    BigtableDataClient failureClient = BigtableDataClient.create(settingsBuilder.build());
+    EnhancedBigtableStub stub = EnhancedBigtableStub.create(settingsBuilder.stubSettings().build());
 
     service.expectations.add(
         RpcExpectation.create()
@@ -928,17 +927,14 @@ public void readRowsAppendsLargeRowsToMidStreamError() throws Exception {
 
     try {
       List<Row> ignored =
-          failureClient
-              .skipLargeRowsCallable()
-              .all()
-              .call(Query.create(TABLE_ID).range("r1", "r5"));
+          stub.deferLargeRowsCallable().all().call(Query.create(TABLE_ID).range("r1", "r5"));
       Truth.assert_().withMessage("Expected InvalidArgumentException").fail();
     } catch (ApiException e) {
       LargeRowException lre = findLargeRowException(e);
       assertThat(lre).isNotNull();
       assertThat(lre.getLargeRowKeys()).containsExactly(ByteString.copyFromUtf8("r2"));
     } finally {
-      failureClient.close();
+      stub.close();
     }
   }