From 9ae5e8418d77742ea8ecc427a616d8596b5d670a Mon Sep 17 00:00:00 2001
From: Teng Zhong <tengzhong@google.com>
Date: Thu, 1 Sep 2022 16:13:53 -0400
Subject: [PATCH] feat: Expose CDC data API settings in
 EnhancedBigtableStubSettings

---
 .../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 49eb79f5ca..f31c5308c5 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
@@ -884,6 +884,20 @@ public UnaryCallSettings.Builder<ReadModifyWriteRow, Row> readModifyWriteRowSett
       return readModifyWriteRowSettings;
     }
 
+    /** 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 731ba7f77e..413bba34d4 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
@@ -634,6 +634,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 =