From 41b07c052e01a97cec1b507599ae138a306c0bc3 Mon Sep 17 00:00:00 2001
From: Israel Herraiz <ihr@google.com>
Date: Fri, 6 Oct 2023 12:25:10 +0200
Subject: [PATCH] Add withAppProfileId to BigtableIO Read and Write

---
 .../beam/sdk/io/gcp/bigtable/BigtableIO.java  | 44 +++++++++++++++++++
 .../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 ef868e8bf7c9..ad978e95016a 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 Read withTableId(String tableId) {
       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 Write withTableId(String tableId) {
       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 971b91d89b74..617d834bcf09 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 void testReadBuildsCorrectly() {
             .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 void testWriteBuildsCorrectly() {
             .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