Add Cloud Spanner Directed Reads support to SpannerIO

This change allows Beam pipeline developers to configure Cloud Spanner DirectedReadOptions on SpannerConfig, as well as on SpannerIO.read(), readAll(), and readChangeStream().

- Adds getDirectedReadOptions() and withDirectedReadOptions(...) overloads (accepting DirectedReadOptions, ValueProvider<DirectedReadOptions>, and String) to SpannerConfig.
- Adds an internal parseDirectedReadOptions helper supporting both standard JSON strings and shorthand <location>:<type> strings (e.g. us-central1:READ_ONLY).
- Threads DirectedReadOptions into SpannerOptions.Builder in SpannerAccessor so that all DatabaseClient instances and ChangeStreamDao queries automatically inherit the directed read configuration.
- Adds delegation methods to SpannerIO.Read, ReadAll, and ReadChangeStream.
- Adds comprehensive unit tests in SpannerAccessorTest and SpannerIOReadChangeStreamTest.

Author: Atharva Moroney

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerAccessor.java

@@ -36,6 +36,7 @@
 import com.google.cloud.spanner.v1.stub.SpannerStubSettings;
 import com.google.spanner.v1.CommitRequest;
 import com.google.spanner.v1.CommitResponse;
+import com.google.spanner.v1.DirectedReadOptions;
 import com.google.spanner.v1.ExecuteSqlRequest;
 import com.google.spanner.v1.PartialResultSet;
 import java.util.HashSet;
@@ -279,6 +280,10 @@ static SpannerOptions buildSpannerOptions(SpannerConfig spannerConfig) {
     if (databaseRole != null && databaseRole.get() != null && !databaseRole.get().isEmpty()) {
       builder.setDatabaseRole(databaseRole.get());
     }
+    ValueProvider<DirectedReadOptions> directedReadOptions = spannerConfig.getDirectedReadOptions();
+    if (directedReadOptions != null && directedReadOptions.get() != null) {
+      builder.setDirectedReadOptions(directedReadOptions.get());
+    }
     ValueProvider<Credentials> credentials = spannerConfig.getCredentials();
     if (credentials != null && credentials.get() != null) {
       builder.setCredentials(credentials.get());

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java

@@ -28,10 +28,15 @@
 import com.google.cloud.spanner.Options.RpcPriority;
 import com.google.cloud.spanner.Spanner;
 import com.google.cloud.spanner.SpannerOptions;
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.util.JsonFormat;
+import com.google.spanner.v1.DirectedReadOptions;
 import java.io.Serializable;
+import java.util.List;
 import org.apache.beam.sdk.options.ValueProvider;
 import org.apache.beam.sdk.transforms.display.DisplayData;
 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.Splitter;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Duration;
@@ -91,6 +96,8 @@ public String getHostValue() {
 
   public abstract @Nullable ValueProvider<String> getDatabaseRole();
 
+  public abstract @Nullable ValueProvider<DirectedReadOptions> getDirectedReadOptions();
+
   public abstract @Nullable ValueProvider<Duration> getPartitionQueryTimeout();
 
   public abstract @Nullable ValueProvider<Duration> getPartitionReadTimeout();
@@ -185,6 +192,8 @@ abstract Builder setExecuteStreamingSqlRetrySettings(
 
     abstract Builder setDatabaseRole(ValueProvider<String> databaseRole);
 
+    abstract Builder setDirectedReadOptions(ValueProvider<DirectedReadOptions> directedReadOptions);
+
     abstract Builder setDataBoostEnabled(ValueProvider<Boolean> dataBoostEnabled);
 
     abstract Builder setPartitionQueryTimeout(ValueProvider<Duration> partitionQueryTimeout);
@@ -335,6 +344,75 @@ public SpannerConfig withDatabaseRole(ValueProvider<String> databaseRole) {
     return toBuilder().setDatabaseRole(databaseRole).build();
   }
 
+  /** Specifies the Cloud Spanner directed read options. */
+  public SpannerConfig withDirectedReadOptions(DirectedReadOptions directedReadOptions) {
+    return withDirectedReadOptions(ValueProvider.StaticValueProvider.of(directedReadOptions));
+  }
+
+  /** Specifies the Cloud Spanner directed read options. */
+  public SpannerConfig withDirectedReadOptions(
+      ValueProvider<DirectedReadOptions> directedReadOptions) {
+    return toBuilder().setDirectedReadOptions(directedReadOptions).build();
+  }
+
+  /** Specifies the Cloud Spanner directed read options from a string representation. */
+  public SpannerConfig withDirectedReadOptions(String directedReadOptions) {
+    if (directedReadOptions == null || directedReadOptions.isEmpty()) {
+      return this;
+    }
+    return withDirectedReadOptions(parseDirectedReadOptions(directedReadOptions));
+  }
+
+  @VisibleForTesting
+  static DirectedReadOptions parseDirectedReadOptions(String directedReadOptions) {
+    if (directedReadOptions == null || directedReadOptions.isEmpty()) {
+      return DirectedReadOptions.getDefaultInstance();
+    }
+    if (!directedReadOptions.trim().startsWith("{")) {
+      List<String> parts = Splitter.on(':').splitToList(directedReadOptions);
+      DirectedReadOptions.ReplicaSelection.Builder selectionBuilder =
+          DirectedReadOptions.ReplicaSelection.newBuilder();
+      if (parts.size() == 2) {
+        selectionBuilder.setLocation(parts.get(0).trim());
+        try {
+          selectionBuilder.setType(
+              DirectedReadOptions.ReplicaSelection.Type.valueOf(parts.get(1).trim().toUpperCase()));
+        } catch (IllegalArgumentException e) {
+          throw new IllegalArgumentException(
+              "Invalid replica type in directed read options: " + parts.get(1), e);
+        }
+      } else if (parts.size() == 1) {
+        try {
+          selectionBuilder.setType(
+              DirectedReadOptions.ReplicaSelection.Type.valueOf(parts.get(0).trim().toUpperCase()));
+        } catch (IllegalArgumentException e) {
+          throw new IllegalArgumentException(
+              "Expected <location>:<type> or valid JSON for directed read options, but got: "
+                  + directedReadOptions,
+              e);
+        }
+      } else {
+        throw new IllegalArgumentException(
+            "Expected <location>:<type> or valid JSON for directed read options, but got: "
+                + directedReadOptions);
+      }
+      return DirectedReadOptions.newBuilder()
+          .setIncludeReplicas(
+              DirectedReadOptions.IncludeReplicas.newBuilder()
+                  .addReplicaSelections(selectionBuilder.build()))
+          .build();
+    }
+
+    DirectedReadOptions.Builder builder = DirectedReadOptions.newBuilder();
+    try {
+      JsonFormat.parser().merge(directedReadOptions, builder);
+      return builder.build();
+    } catch (InvalidProtocolBufferException e) {
+      throw new IllegalArgumentException(
+          "Failed to parse DirectedReadOptions JSON: " + directedReadOptions, e);
+    }
+  }
+
   /** Specifies if the pipeline has to be run on the independent compute resource. */
   public SpannerConfig withDataBoostEnabled(ValueProvider<Boolean> dataBoostEnabled) {
     return toBuilder().setDataBoostEnabled(dataBoostEnabled).build();

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java

@@ -61,6 +61,7 @@
 import com.google.cloud.spanner.TimestampBound;
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
+import com.google.spanner.v1.DirectedReadOptions;
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
@@ -638,6 +639,24 @@ public ReadAll withExperimentalHost(String experimentalHost) {
       return withExperimentalHost(ValueProvider.StaticValueProvider.of(experimentalHost));
     }
 
+    /** Specifies the directed read options for Cloud Spanner. */
+    public ReadAll withDirectedReadOptions(DirectedReadOptions directedReadOptions) {
+      SpannerConfig config = getSpannerConfig();
+      return withSpannerConfig(config.withDirectedReadOptions(directedReadOptions));
+    }
+
+    /** Specifies the directed read options for Cloud Spanner. */
+    public ReadAll withDirectedReadOptions(ValueProvider<DirectedReadOptions> directedReadOptions) {
+      SpannerConfig config = getSpannerConfig();
+      return withSpannerConfig(config.withDirectedReadOptions(directedReadOptions));
+    }
+
+    /** Specifies the directed read options for Cloud Spanner from a string representation. */
+    public ReadAll withDirectedReadOptions(String directedReadOptions) {
+      SpannerConfig config = getSpannerConfig();
+      return withSpannerConfig(config.withDirectedReadOptions(directedReadOptions));
+    }
+
     /**
      * Specifies whether to use plaintext channel.
      *
@@ -927,6 +946,24 @@ public Read withExperimentalHost(String experimentalHost) {
       return withExperimentalHost(ValueProvider.StaticValueProvider.of(experimentalHost));
     }
 
+    /** Specifies the directed read options for Cloud Spanner. */
+    public Read withDirectedReadOptions(DirectedReadOptions directedReadOptions) {
+      SpannerConfig config = getSpannerConfig();
+      return withSpannerConfig(config.withDirectedReadOptions(directedReadOptions));
+    }
+
+    /** Specifies the directed read options for Cloud Spanner. */
+    public Read withDirectedReadOptions(ValueProvider<DirectedReadOptions> directedReadOptions) {
+      SpannerConfig config = getSpannerConfig();
+      return withSpannerConfig(config.withDirectedReadOptions(directedReadOptions));
+    }
+
+    /** Specifies the directed read options for Cloud Spanner from a string representation. */
+    public Read withDirectedReadOptions(String directedReadOptions) {
+      SpannerConfig config = getSpannerConfig();
+      return withSpannerConfig(config.withDirectedReadOptions(directedReadOptions));
+    }
+
     /**
      * Specifies whether to use plaintext channel.
      *
@@ -2052,6 +2089,28 @@ public ReadChangeStream withExperimentalHost(String experimentalHost) {
       return withExperimentalHost(ValueProvider.StaticValueProvider.of(experimentalHost));
     }
 
+    /** Specifies the directed read options for change stream queries. */
+    public ReadChangeStream withDirectedReadOptions(DirectedReadOptions directedReadOptions) {
+      SpannerConfig config = getSpannerConfig();
+      return withSpannerConfig(config.withDirectedReadOptions(directedReadOptions));
+    }
+
+    /** Specifies the directed read options for change stream queries. */
+    public ReadChangeStream withDirectedReadOptions(
+        ValueProvider<DirectedReadOptions> directedReadOptions) {
+      SpannerConfig config = getSpannerConfig();
+      return withSpannerConfig(config.withDirectedReadOptions(directedReadOptions));
+    }
+
+    /**
+     * Specifies the directed read options for change stream queries from a string representation
+     * (e.g., JSON string or "us-central1:READ_ONLY").
+     */
+    public ReadChangeStream withDirectedReadOptions(String directedReadOptions) {
+      SpannerConfig config = getSpannerConfig();
+      return withSpannerConfig(config.withDirectedReadOptions(directedReadOptions));
+    }
+
     /**
      * Specifies whether to use plaintext channel.
      *

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerAccessorTest.java

@@ -24,6 +24,7 @@
 
 import com.google.cloud.spanner.DatabaseId;
 import com.google.cloud.spanner.SpannerOptions;
+import com.google.spanner.v1.DirectedReadOptions;
 import org.apache.beam.sdk.extensions.gcp.auth.TestCredential;
 import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
 import org.junit.Before;
@@ -251,4 +252,78 @@ public void testBuildSpannerOptionsWithCustomHost() {
     SpannerOptions options = SpannerAccessor.buildSpannerOptions(config1);
     assertEquals(host, options.getHost());
   }
+
+  @Test
+  public void testBuildSpannerOptionsWithDirectedReadOptions() {
+    DirectedReadOptions directedReadOptions =
+        DirectedReadOptions.newBuilder()
+            .setIncludeReplicas(
+                DirectedReadOptions.IncludeReplicas.newBuilder()
+                    .addReplicaSelections(
+                        DirectedReadOptions.ReplicaSelection.newBuilder()
+                            .setLocation("us-central1")
+                            .setType(DirectedReadOptions.ReplicaSelection.Type.READ_ONLY)))
+            .build();
+    SpannerConfig config1 =
+        SpannerConfig.create()
+            .toBuilder()
+            .setServiceFactory(serviceFactory)
+            .setDirectedReadOptions(StaticValueProvider.of(directedReadOptions))
+            .setProjectId(StaticValueProvider.of("project"))
+            .setInstanceId(StaticValueProvider.of("test1"))
+            .setDatabaseId(StaticValueProvider.of("test1"))
+            .build();
+
+    SpannerOptions options = SpannerAccessor.buildSpannerOptions(config1);
+    assertEquals(directedReadOptions, options.getDirectedReadOptions());
+  }
+
+  @Test
+  public void testBuildSpannerOptionsWithDirectedReadOptionsString() {
+    String directedReadString = "us-central1:READ_ONLY";
+    SpannerConfig config1 =
+        SpannerConfig.create()
+            .withServiceFactory(serviceFactory)
+            .withProjectId("project")
+            .withInstanceId("test1")
+            .withDatabaseId("test1")
+            .withDirectedReadOptions(directedReadString);
+
+    SpannerOptions options = SpannerAccessor.buildSpannerOptions(config1);
+    assertEquals(
+        DirectedReadOptions.ReplicaSelection.Type.READ_ONLY,
+        options.getDirectedReadOptions().getIncludeReplicas().getReplicaSelections(0).getType());
+    assertEquals(
+        "us-central1",
+        options
+            .getDirectedReadOptions()
+            .getIncludeReplicas()
+            .getReplicaSelections(0)
+            .getLocation());
+  }
+
+  @Test
+  public void testBuildSpannerOptionsWithDirectedReadOptionsJson() {
+    String jsonString =
+        "{\"includeReplicas\":{\"replicaSelections\":[{\"location\":\"us-east1\",\"type\":\"READ_WRITE\"}]}}";
+    SpannerConfig config1 =
+        SpannerConfig.create()
+            .withServiceFactory(serviceFactory)
+            .withProjectId("project")
+            .withInstanceId("test1")
+            .withDatabaseId("test1")
+            .withDirectedReadOptions(jsonString);
+
+    SpannerOptions options = SpannerAccessor.buildSpannerOptions(config1);
+    assertEquals(
+        DirectedReadOptions.ReplicaSelection.Type.READ_WRITE,
+        options.getDirectedReadOptions().getIncludeReplicas().getReplicaSelections(0).getType());
+    assertEquals(
+        "us-east1",
+        options
+            .getDirectedReadOptions()
+            .getIncludeReplicas()
+            .getReplicaSelections(0)
+            .getLocation());
+  }
 }

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOReadChangeStreamTest.java

@@ -159,6 +159,21 @@ public void testWithDefaultCredential() {
           defaultCredential, changeStreamSpannerConfigWithCredential.getCredentials().get());
       assertEquals(defaultCredential, metadataSpannerConfigWithCredential.getCredentials().get());
     }
+
+    @Test
+    public void testSetDirectedReadOptions() {
+      String directedReadString = "us-central1:READ_ONLY";
+      readChangeStream = readChangeStream.withDirectedReadOptions(directedReadString);
+      SpannerConfig changeStreamSpannerConfig = readChangeStream.buildChangeStreamSpannerConfig();
+      assertEquals(
+          "us-central1",
+          changeStreamSpannerConfig
+              .getDirectedReadOptions()
+              .get()
+              .getIncludeReplicas()
+              .getReplicaSelections(0)
+              .getLocation());
+    }
   }
 
   /** Parameterized tests for Dialect and Partition Mode combinations. */