From 31c1c8d07923119a8b3f3782ee9c11a957e38758 Mon Sep 17 00:00:00 2001
From: tengzhonger <109308630+tengzhonger@users.noreply.github.com>
Date: Thu, 1 Sep 2022 16:16:22 -0400
Subject: [PATCH] feat: Expose CDC data API settings in
 EnhancedBigtableStubSettings (#1376)

Co-authored-by: Teng Zhong <tengzhong@google.com>
---
 .../v2/stub/EnhancedBigtableStubSettings.java | 14 ++++
 .../EnhancedBigtableStubSettingsTest.java     | 75 +++++++++++++++++++
 2 files changed, 89 insertions(+)

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java
index f435a5d685..9c74642a71 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java
+++ b/google-cloud-bigtable/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/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java
index 92924aa7a6..c25c231f85 100644
--- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java
+++ b/google-cloud-bigtable/src/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 =