[beam] branch master updated: Add withAppProfileId to BigtableIO Read and Write (#28864)

[johncasey]

This is an automated email from the ASF dual-hosted git repository.

johncasey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new 4e6f10b5be5 Add withAppProfileId to BigtableIO Read and Write (#28864)
4e6f10b5be5 is described below

commit 4e6f10b5be51f90f9410b2a16652fb2c71625b77
Author: Israel Herraiz <i...@google.com>
AuthorDate: Wed Oct 18 16:59:35 2023 +0200

    Add withAppProfileId to BigtableIO Read and Write (#28864)
---
 .../beam/sdk/io/gcp/bigtable/BigtableIO.java       | 44 ++++++++++++++++++++++
 .../beam/sdk/io/gcp/bigtable/BigtableIOTest.java   |  6 ++-
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git 
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java
 
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java
index ef868e8bf7c..ad978e95016 100644
--- 
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java
+++ 
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java
@@ -457,6 +457,25 @@ public class BigtableIO {
       return withTableId(StaticValueProvider.of(tableId));
     }
 
+    /**
+     * Returns a new {@link BigtableIO.Read} that will read using the 
specified app profile id.
+     *
+     * <p>Does not modify this object.
+     */
+    public Read withAppProfileId(ValueProvider<String> appProfileId) {
+      BigtableConfig config = getBigtableConfig();
+      return 
toBuilder().setBigtableConfig(config.withAppProfileId(appProfileId)).build();
+    }
+
+    /**
+     * Returns a new {@link BigtableIO.Read} that will read using the 
specified app profile id.
+     *
+     * <p>Does not modify this object.
+     */
+    public Read withAppProfileId(String appProfileId) {
+      return withAppProfileId(StaticValueProvider.of(appProfileId));
+    }
+
     /**
      * WARNING: Should be used only to specify additional parameters for 
connection to the Cloud
      * Bigtable, instanceId and projectId should be provided over {@link 
#withInstanceId} and {@link
@@ -837,6 +856,31 @@ public class BigtableIO {
       return withTableId(StaticValueProvider.of(tableId));
     }
 
+    /**
+     * Returns a new {@link BigtableIO.Write} that will write using the 
specified app profile id.
+     *
+     * <p>Remember that in order to use single-row transactions, this must use 
a single-cluster
+     * routing policy.
+     *
+     * <p>Does not modify this object.
+     */
+    public Write withAppProfileId(ValueProvider<String> appProfileId) {
+      BigtableConfig config = getBigtableConfig();
+      return 
toBuilder().setBigtableConfig(config.withAppProfileId(appProfileId)).build();
+    }
+
+    /**
+     * Returns a new {@link BigtableIO.Write} that will write using the 
specified app profile id.
+     *
+     * <p>Remember that in order to use single-row transactions, this must use 
a single-cluster
+     * routing policy.
+     *
+     * <p>Does not modify this object.
+     */
+    public Write withAppProfileId(String appProfileId) {
+      return withAppProfileId(StaticValueProvider.of(appProfileId));
+    }
+
     /**
      * WARNING: Should be used only to specify additional parameters for 
connection to the Cloud
      * Bigtable, instanceId and projectId should be provided over {@link 
#withInstanceId} and {@link
diff --git 
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIOTest.java
 
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIOTest.java
index f6224cc24b2..714dc9f8619 100644
--- 
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIOTest.java
+++ 
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIOTest.java
@@ -219,11 +219,13 @@ public class BigtableIOTest {
             .withTableId("table")
             .withInstanceId("instance")
             .withProjectId("project")
+            .withAppProfileId("app-profile")
             .withBigtableOptionsConfigurator(PORT_CONFIGURATOR);
     assertEquals("options_project", read.getBigtableOptions().getProjectId());
     assertEquals("options_instance", 
read.getBigtableOptions().getInstanceId());
     assertEquals("instance", read.getBigtableConfig().getInstanceId().get());
     assertEquals("project", read.getBigtableConfig().getProjectId().get());
+    assertEquals("app-profile", 
read.getBigtableConfig().getAppProfileId().get());
     assertEquals("table", read.getTableId());
     assertEquals(PORT_CONFIGURATOR, 
read.getBigtableConfig().getBigtableOptionsConfigurator());
   }
@@ -373,12 +375,14 @@ public class BigtableIOTest {
             .withBigtableOptions(BIGTABLE_OPTIONS)
             .withTableId("table")
             .withInstanceId("instance")
-            .withProjectId("project");
+            .withProjectId("project")
+            .withAppProfileId("app-profile");
     assertEquals("table", write.getBigtableWriteOptions().getTableId().get());
     assertEquals("options_project", write.getBigtableOptions().getProjectId());
     assertEquals("options_instance", 
write.getBigtableOptions().getInstanceId());
     assertEquals("instance", write.getBigtableConfig().getInstanceId().get());
     assertEquals("project", write.getBigtableConfig().getProjectId().get());
+    assertEquals("app-profile", 
write.getBigtableConfig().getAppProfileId().get());
   }
 
   @Test