From 49c99ee030dd591d373690093fb4f59854a059e6 Mon Sep 17 00:00:00 2001
From: Mairbek Khadikov <mairbek@google.com>
Date: Tue, 26 Sep 2017 16:59:30 -0700
Subject: [PATCH] Allows to set a Cloud Spanner host.
 https://batch-spanner.googleapis.com/ is set as a default host name.

---
 .../sdk/io/gcp/spanner/SpannerConfig.java     | 16 ++++++++++++-
 .../beam/sdk/io/gcp/spanner/SpannerIO.java    | 24 +++++++++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java
index dc0bab8e0d32..9be641fbf884 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java
@@ -37,6 +37,8 @@
 public abstract class SpannerConfig implements Serializable {
   // A common user agent token that indicates that this request was originated from Apache Beam.
   private static final String USER_AGENT_PREFIX = "Apache_Beam_Java";
+  // A default host name for batch traffic.
+  private static final String DEFAULT_HOST = "https://batch-spanner.googleapis.com/";
 
   @Nullable
   abstract ValueProvider<String> getProjectId();
@@ -47,6 +49,9 @@ public abstract class SpannerConfig implements Serializable {
   @Nullable
   abstract ValueProvider<String> getDatabaseId();
 
+  @Nullable
+  abstract String getHost();
+
   @Nullable
   @VisibleForTesting
   abstract ServiceFactory<Spanner, SpannerOptions> getServiceFactory();
@@ -54,7 +59,7 @@ public abstract class SpannerConfig implements Serializable {
   abstract Builder toBuilder();
 
   public static SpannerConfig create() {
-    return builder().build();
+    return builder().setHost(DEFAULT_HOST).build();
   }
 
   static Builder builder() {
@@ -93,6 +98,8 @@ public abstract static class Builder {
 
     abstract Builder setDatabaseId(ValueProvider<String> databaseId);
 
+    abstract Builder setHost(String host);
+
     abstract Builder setServiceFactory(ServiceFactory<Spanner, SpannerOptions> serviceFactory);
 
     public abstract SpannerConfig build();
@@ -122,6 +129,10 @@ public SpannerConfig withDatabaseId(String databaseId) {
     return withDatabaseId(ValueProvider.StaticValueProvider.of(databaseId));
   }
 
+  public SpannerConfig withHost(String host) {
+    return toBuilder().setHost(host).build();
+  }
+
   @VisibleForTesting
   SpannerConfig withServiceFactory(ServiceFactory<Spanner, SpannerOptions> serviceFactory) {
     return toBuilder().setServiceFactory(serviceFactory).build();
@@ -135,6 +146,9 @@ public SpannerAccessor connectToSpanner() {
     if (getServiceFactory() != null) {
       builder.setServiceFactory(this.getServiceFactory());
     }
+    if (getHost() != null) {
+      builder.setHost(getHost());
+    }
     ReleaseInfo releaseInfo = ReleaseInfo.getReleaseInfo();
     builder.setUserAgentPrefix(USER_AGENT_PREFIX + "/" + releaseInfo.getVersion());
     SpannerOptions options = builder.build();
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
index 54759fb69105..bd0e6ee3171c 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
@@ -269,6 +269,12 @@ public ReadAll withDatabaseId(String databaseId) {
       return withDatabaseId(ValueProvider.StaticValueProvider.of(databaseId));
     }
 
+    /** Specifies the Cloud Spanner host. */
+    public ReadAll witHost(String host) {
+      SpannerConfig config = getSpannerConfig();
+      return withSpannerConfig(config.withHost(host));
+    }
+
     /** Specifies the Cloud Spanner database. */
     public ReadAll withDatabaseId(ValueProvider<String> databaseId) {
       SpannerConfig config = getSpannerConfig();
@@ -380,6 +386,12 @@ public Read withDatabaseId(ValueProvider<String> databaseId) {
       return withSpannerConfig(config.withDatabaseId(databaseId));
     }
 
+    /** Specifies the Cloud Spanner host. */
+    public Read witHost(String host) {
+      SpannerConfig config = getSpannerConfig();
+      return withSpannerConfig(config.withHost(host));
+    }
+
     @VisibleForTesting
     Read withServiceFactory(ServiceFactory<Spanner, SpannerOptions> serviceFactory) {
       SpannerConfig config = getSpannerConfig();
@@ -533,6 +545,12 @@ public CreateTransaction withDatabaseId(ValueProvider<String> databaseId) {
       return withSpannerConfig(config.withDatabaseId(databaseId));
     }
 
+    /** Specifies the Cloud Spanner host. */
+    public CreateTransaction witHost(String host) {
+      SpannerConfig config = getSpannerConfig();
+      return withSpannerConfig(config.withHost(host));
+    }
+
     @VisibleForTesting
     CreateTransaction withServiceFactory(
         ServiceFactory<Spanner, SpannerOptions> serviceFactory) {
@@ -619,6 +637,12 @@ public Write withDatabaseId(ValueProvider<String> databaseId) {
       return withSpannerConfig(config.withDatabaseId(databaseId));
     }
 
+    /** Specifies the Cloud Spanner host. */
+    public Write witHost(String host) {
+      SpannerConfig config = getSpannerConfig();
+      return withSpannerConfig(config.withHost(host));
+    }
+
     @VisibleForTesting
     Write withServiceFactory(ServiceFactory<Spanner, SpannerOptions> serviceFactory) {
       SpannerConfig config = getSpannerConfig();