Use only ValueProviders in SpannerConfig (apache#24156)

andreigurau · Nov 17, 2022 · 2bbaf35 · 2bbaf35
1 parent a4c7cbc
commit 2bbaf35
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 8 deletions.
diff --git a/...ogle-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerAccessor.java b/...ogle-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerAccessor.java
@@ -185,9 +185,9 @@ private static SpannerAccessor createAndConnect(SpannerConfig spannerConfig) {
     }
     String userAgentString = USER_AGENT_PREFIX + "/" + ReleaseInfo.getReleaseInfo().getVersion();
     builder.setHeaderProvider(FixedHeaderProvider.create("user-agent", userAgentString));
-    String databaseRole = spannerConfig.getDatabaseRole();
-    if (databaseRole != null && !databaseRole.isEmpty()) {
-      builder.setDatabaseRole(databaseRole);
+    ValueProvider<String> databaseRole = spannerConfig.getDatabaseRole();
+    if (databaseRole != null && databaseRole.get() != null && !databaseRole.get().isEmpty()) {
+      builder.setDatabaseRole(databaseRole.get());
     }
     SpannerOptions options = builder.build();
 

diff --git a/...google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java b/...google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java
@@ -73,7 +73,7 @@ public abstract class SpannerConfig implements Serializable {
 
   public abstract @Nullable ValueProvider<RpcPriority> getRpcPriority();
 
-  public abstract @Nullable String getDatabaseRole();
+  public abstract @Nullable ValueProvider<String> getDatabaseRole();
 
   @VisibleForTesting
   abstract @Nullable ServiceFactory<Spanner, SpannerOptions> getServiceFactory();
@@ -147,7 +147,7 @@ abstract Builder setExecuteStreamingSqlRetrySettings(
 
     abstract Builder setRpcPriority(ValueProvider<RpcPriority> rpcPriority);
 
-    abstract Builder setDatabaseRole(String databaseRole);
+    abstract Builder setDatabaseRole(ValueProvider<String> databaseRole);
 
     public abstract SpannerConfig build();
   }
@@ -262,7 +262,7 @@ public SpannerConfig withRpcPriority(ValueProvider<RpcPriority> rpcPriority) {
   }
 
   /** Specifies the Cloud Spanner database role. */
-  public SpannerConfig withDatabaseRole(String databaseRole) {
+  public SpannerConfig withDatabaseRole(ValueProvider<String> databaseRole) {
     return toBuilder().setDatabaseRole(databaseRole).build();
   }
 }
diff --git a/...-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerAccessorTest.java b/...-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerAccessorTest.java
@@ -109,7 +109,7 @@ public void testCreateWithValidDatabaseRole() {
             .setProjectId(StaticValueProvider.of("project"))
             .setInstanceId(StaticValueProvider.of("test1"))
             .setDatabaseId(StaticValueProvider.of("test1"))
-            .setDatabaseRole("test-role")
+            .setDatabaseRole(StaticValueProvider.of("test-role"))
             .build();
 
     SpannerAccessor acc1 = SpannerAccessor.getOrCreate(config1);
@@ -130,7 +130,7 @@ public void testCreateWithEmptyDatabaseRole() {
             .setProjectId(StaticValueProvider.of("project"))
             .setInstanceId(StaticValueProvider.of("test1"))
             .setDatabaseId(StaticValueProvider.of("test1"))
-            .setDatabaseRole("")
+            .setDatabaseRole(StaticValueProvider.of(""))
             .build();
 
     SpannerAccessor acc1 = SpannerAccessor.getOrCreate(config1);