Sync SpannerConfig with Beam

GoogleCloudPlatform · Mar 13, 2024 · 447cfaa · 447cfaa
1 parent 0c1d028
commit 447cfaa
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/v1/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java b/v1/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java
@@ -21,6 +21,7 @@
 
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.api.gax.rpc.StatusCode.Code;
+import com.google.auth.Credentials;
 import com.google.auto.value.AutoValue;
 import com.google.cloud.ServiceFactory;
 import com.google.cloud.spanner.Options.RpcPriority;
@@ -86,6 +87,8 @@ public abstract class SpannerConfig implements Serializable {
 
   public abstract @Nullable ValueProvider<Boolean> getDataBoostEnabled();
 
+  public abstract @Nullable ValueProvider<Credentials> getCredentials();
+
   abstract Builder toBuilder();
 
   public static SpannerConfig create() {
@@ -164,6 +167,8 @@ abstract Builder setExecuteStreamingSqlRetrySettings(
 
     abstract Builder setPartitionReadTimeout(ValueProvider<Duration> partitionReadTimeout);
 
+    abstract Builder setCredentials(ValueProvider<Credentials> credentials);
+
     public abstract SpannerConfig build();
   }
 
@@ -305,4 +310,14 @@ public SpannerConfig withPartitionReadTimeout(Duration partitionReadTimeout) {
   public SpannerConfig withPartitionReadTimeout(ValueProvider<Duration> partitionReadTimeout) {
     return toBuilder().setPartitionReadTimeout(partitionReadTimeout).build();
   }
+
+  /** Specifies the credentials. */
+  public SpannerConfig withCredentials(Credentials credentials) {
+    return withCredentials(ValueProvider.StaticValueProvider.of(credentials));
+  }
+
+  /** Specifies the credentials. */
+  public SpannerConfig withCredentials(ValueProvider<Credentials> credentials) {
+    return toBuilder().setCredentials(credentials).build();
+  }
 }