feat: Expose CDC data API settings in EnhancedBigtableStubSettings (g…

…oogleapis#1376) Co-authored-by: Teng Zhong <[email protected]>
tengzhonger · Sep 2, 2022 · 31c1c8d · 31c1c8d
1 parent 7d58536
commit 31c1c8d
Show file tree

Hide file tree

Showing 2 changed files with 89 additions and 0 deletions.
diff --git a/...le/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java b/...le/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java
@@ -925,6 +925,20 @@ public UnaryCallSettings.Builder<PingAndWarmRequest, Void> pingAndWarmSettings()
       return pingAndWarmSettings;
     }
 
+    /** Returns the builder for the settings used for calls to ReadChangeStream. */
+    public ServerStreamingCallSettings.Builder<ReadChangeStreamQuery, ChangeStreamRecord>
+        readChangeStreamSettings() {
+      return readChangeStreamSettings;
+    }
+
+    /**
+     * Returns the builder for the settings used for calls to GenerateInitialChangeStreamPartitions.
+     */
+    public ServerStreamingCallSettings.Builder<String, ByteStringRange>
+        generateInitialChangeStreamPartitionsSettings() {
+      return generateInitialChangeStreamPartitionsSettings;
+    }
+
     @SuppressWarnings("unchecked")
     public EnhancedBigtableStubSettings build() {
       Preconditions.checkState(projectId != null, "Project id must be set");

diff --git a/...rc/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java b/...rc/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java
@@ -635,6 +635,81 @@ public void checkAndMutateRowSettingsAreNotLostTest() {
         .isEqualTo(retrySettings);
   }
 
+  @Test
+  public void generateInitialChangeStreamPartitionsSettingsAreNotLostTest() {
+    String dummyProjectId = "my-project";
+    String dummyInstanceId = "my-instance";
+
+    EnhancedBigtableStubSettings.Builder builder =
+        EnhancedBigtableStubSettings.newBuilder()
+            .setProjectId(dummyProjectId)
+            .setInstanceId(dummyInstanceId);
+
+    RetrySettings retrySettings = RetrySettings.newBuilder().build();
+    builder
+        .generateInitialChangeStreamPartitionsSettings()
+        .setRetryableCodes(Code.ABORTED, Code.DEADLINE_EXCEEDED)
+        .setRetrySettings(retrySettings)
+        .build();
+
+    assertThat(builder.generateInitialChangeStreamPartitionsSettings().getRetryableCodes())
+        .containsAtLeast(Code.ABORTED, Code.DEADLINE_EXCEEDED);
+    assertThat(builder.generateInitialChangeStreamPartitionsSettings().getRetrySettings())
+        .isEqualTo(retrySettings);
+
+    assertThat(builder.build().generateInitialChangeStreamPartitionsSettings().getRetryableCodes())
+        .containsAtLeast(Code.ABORTED, Code.DEADLINE_EXCEEDED);
+    assertThat(builder.build().generateInitialChangeStreamPartitionsSettings().getRetrySettings())
+        .isEqualTo(retrySettings);
+
+    assertThat(
+            builder
+                .build()
+                .toBuilder()
+                .generateInitialChangeStreamPartitionsSettings()
+                .getRetryableCodes())
+        .containsAtLeast(Code.ABORTED, Code.DEADLINE_EXCEEDED);
+    assertThat(
+            builder
+                .build()
+                .toBuilder()
+                .generateInitialChangeStreamPartitionsSettings()
+                .getRetrySettings())
+        .isEqualTo(retrySettings);
+  }
+
+  @Test
+  public void readChangeStreamSettingsAreNotLostTest() {
+    String dummyProjectId = "my-project";
+    String dummyInstanceId = "my-instance";
+
+    EnhancedBigtableStubSettings.Builder builder =
+        EnhancedBigtableStubSettings.newBuilder()
+            .setProjectId(dummyProjectId)
+            .setInstanceId(dummyInstanceId);
+
+    RetrySettings retrySettings = RetrySettings.newBuilder().build();
+    builder
+        .readChangeStreamSettings()
+        .setRetryableCodes(Code.ABORTED, Code.DEADLINE_EXCEEDED)
+        .setRetrySettings(retrySettings)
+        .build();
+
+    assertThat(builder.readChangeStreamSettings().getRetryableCodes())
+        .containsAtLeast(Code.ABORTED, Code.DEADLINE_EXCEEDED);
+    assertThat(builder.readChangeStreamSettings().getRetrySettings()).isEqualTo(retrySettings);
+
+    assertThat(builder.build().readChangeStreamSettings().getRetryableCodes())
+        .containsAtLeast(Code.ABORTED, Code.DEADLINE_EXCEEDED);
+    assertThat(builder.build().readChangeStreamSettings().getRetrySettings())
+        .isEqualTo(retrySettings);
+
+    assertThat(builder.build().toBuilder().readChangeStreamSettings().getRetryableCodes())
+        .containsAtLeast(Code.ABORTED, Code.DEADLINE_EXCEEDED);
+    assertThat(builder.build().toBuilder().readChangeStreamSettings().getRetrySettings())
+        .isEqualTo(retrySettings);
+  }
+
   @Test
   public void checkAndMutateRowSettingsAreSane() {
     UnaryCallSettings.Builder<ConditionalRowMutation, Boolean> builder =