(beam) branch master updated: Bigtable: use InstanceBuilder to dynamic load override class (#30542)

[yhu]

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

yhu 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 f0f1cd37dd1 Bigtable: use InstanceBuilder to dynamic load override 
class (#30542)
f0f1cd37dd1 is described below

commit f0f1cd37dd11ba15ee8f88b5e0b81fecdcd3b1b0
Author: Mattie Fu <matti...@google.com>
AuthorDate: Wed Mar 6 12:18:28 2024 -0500

    Bigtable: use InstanceBuilder to dynamic load override class (#30542)
    
    * Bigtable: use InstanceBuilder to dynamic load override class
    
    * remove unnecessary check
---
 .../io/gcp/bigtable/BigtableConfigTranslator.java  | 32 ++++++++++------------
 .../gcp/bigtable/BigtableConfigTranslatorTest.java |  2 +-
 2 files changed, 16 insertions(+), 18 deletions(-)

diff --git 
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableConfigTranslator.java
 
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableConfigTranslator.java
index fcb05cd8a9d..9fa68bdd018 100644
--- 
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableConfigTranslator.java
+++ 
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableConfigTranslator.java
@@ -47,6 +47,8 @@ import org.apache.beam.sdk.options.ExperimentalOptions;
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.options.ValueProvider;
+import org.apache.beam.sdk.util.InstanceBuilder;
+import org.apache.beam.sdk.values.TypeDescriptor;
 import 
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings;
 import 
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
@@ -161,7 +163,6 @@ class BigtableConfigTranslator {
     return configureSettingsOverride(overrideClassName, dataBuilder, 
pipelineOptions);
   }
 
-  @SuppressWarnings("unchecked")
   private static BigtableDataSettings.Builder configureSettingsOverride(
       @Nullable String override,
       BigtableDataSettings.Builder dataBuilder,
@@ -169,25 +170,22 @@ class BigtableConfigTranslator {
     if (override == null) {
       return dataBuilder;
     }
-    Object object;
+
+    BiFunction<BigtableDataSettings.Builder, PipelineOptions, 
BigtableDataSettings.Builder>
+        overrideFunction;
     try {
-      object = Class.forName(override).getConstructor().newInstance();
-    } catch (Exception e) {
+      overrideFunction =
+          InstanceBuilder.ofType(
+                  new TypeDescriptor<
+                      BiFunction<
+                          BigtableDataSettings.Builder,
+                          PipelineOptions,
+                          BigtableDataSettings.Builder>>() {})
+              .fromClassName(override)
+              .build();
+    } catch (ClassNotFoundException e) {
       throw new IllegalArgumentException("Failed to load class override " + 
override, e);
     }
-    if (!(object instanceof BiFunction<?, ?, ?>)) {
-      throw new IllegalArgumentException(
-          "Incorrect override class type for "
-              + override
-              + ", override class need to be a subclass of "
-              + "BiFunction<BigtableDataSettings.Builder, PipelineOptions, 
BigtableDataSettings.Builder>. Actual type is "
-              + object.getClass());
-    }
-    BiFunction<BigtableDataSettings.Builder, PipelineOptions, 
BigtableDataSettings.Builder>
-        overrideFunction =
-            (BiFunction<
-                    BigtableDataSettings.Builder, PipelineOptions, 
BigtableDataSettings.Builder>)
-                object;
     try {
       return overrideFunction.apply(dataBuilder, pipelineOptions);
     } catch (Exception e) {
diff --git 
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableConfigTranslatorTest.java
 
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableConfigTranslatorTest.java
index db64e11fa04..802f04506ef 100644
--- 
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableConfigTranslatorTest.java
+++ 
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableConfigTranslatorTest.java
@@ -511,7 +511,7 @@ public class BigtableConfigTranslatorTest {
 
     BigtableConfig finalConfig = config;
     assertThrows(
-        IllegalArgumentException.class,
+        IllegalStateException.class,
         () -> BigtableConfigTranslator.translateToVeneerSettings(finalConfig, 
pipelineOptions));
   }