From 6b9513f95d261b09f0aa003c08ecb65e09ed79c7 Mon Sep 17 00:00:00 2001
From: Rajat Bhatta <93644539+rajatbhatta@users.noreply.github.com>
Date: Fri, 19 Jan 2024 19:13:20 +0800
Subject: [PATCH] sample: add support for Directed Read options (#2394)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* sample: add support for Directed Read options

* add additional assertion on output from sample

* fix nit

Co-authored-by: Knut Olav Løite <koloite@gmail.com>

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* feat: update setter

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* feat: update comment

* feat: update comment

* feat: add additional comments

* feat: lint fix

---------

Co-authored-by: Knut Olav Løite <koloite@gmail.com>
Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
Co-authored-by: Sri Harsha CH <57220027+harshachinta@users.noreply.github.com>
Co-authored-by: Sri Harsha CH <sriharshach@google.com>
---
 README.md                                     |   1 +
 .../example/spanner/DirectedReadSample.java   | 114 ++++++++++++++++++
 .../example/spanner/DirectedReadSampleIT.java | 105 ++++++++++++++++
 3 files changed, 220 insertions(+)
 create mode 100644 samples/snippets/src/main/java/com/example/spanner/DirectedReadSample.java
 create mode 100644 samples/snippets/src/test/java/com/example/spanner/DirectedReadSampleIT.java

diff --git a/README.md b/README.md
index 503eb3677d5..4784fc5caeb 100644
--- a/README.md
+++ b/README.md
@@ -285,6 +285,7 @@ Samples are in the [`samples/`](https://github.com/googleapis/java-spanner/tree/
 | Custom Timeout And Retry Settings Example | [source code](https://github.com/googleapis/java-spanner/blob/main/samples/snippets/src/main/java/com/example/spanner/CustomTimeoutAndRetrySettingsExample.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-spanner&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/spanner/CustomTimeoutAndRetrySettingsExample.java) |
 | Delete Instance Config Sample | [source code](https://github.com/googleapis/java-spanner/blob/main/samples/snippets/src/main/java/com/example/spanner/DeleteInstanceConfigSample.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-spanner&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/spanner/DeleteInstanceConfigSample.java) |
 | Delete Using Dml Returning Sample | [source code](https://github.com/googleapis/java-spanner/blob/main/samples/snippets/src/main/java/com/example/spanner/DeleteUsingDmlReturningSample.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-spanner&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/spanner/DeleteUsingDmlReturningSample.java) |
+| Directed Read Sample | [source code](https://github.com/googleapis/java-spanner/blob/main/samples/snippets/src/main/java/com/example/spanner/DirectedReadSample.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-spanner&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/spanner/DirectedReadSample.java) |
 | Drop Foreign Key Constraint Delete Cascade Sample | [source code](https://github.com/googleapis/java-spanner/blob/main/samples/snippets/src/main/java/com/example/spanner/DropForeignKeyConstraintDeleteCascadeSample.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-spanner&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/spanner/DropForeignKeyConstraintDeleteCascadeSample.java) |
 | Drop Sequence Sample | [source code](https://github.com/googleapis/java-spanner/blob/main/samples/snippets/src/main/java/com/example/spanner/DropSequenceSample.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-spanner&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/spanner/DropSequenceSample.java) |
 | Enable Fine Grained Access | [source code](https://github.com/googleapis/java-spanner/blob/main/samples/snippets/src/main/java/com/example/spanner/EnableFineGrainedAccess.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-spanner&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/spanner/EnableFineGrainedAccess.java) |
diff --git a/samples/snippets/src/main/java/com/example/spanner/DirectedReadSample.java b/samples/snippets/src/main/java/com/example/spanner/DirectedReadSample.java
new file mode 100644
index 00000000000..141d9e28244
--- /dev/null
+++ b/samples/snippets/src/main/java/com/example/spanner/DirectedReadSample.java
@@ -0,0 +1,114 @@
+/*
+ * Copyright 2023 Google Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.example.spanner;
+
+// [START spanner_directed_read]
+import com.google.cloud.spanner.DatabaseClient;
+import com.google.cloud.spanner.DatabaseId;
+import com.google.cloud.spanner.Options;
+import com.google.cloud.spanner.ResultSet;
+import com.google.cloud.spanner.Spanner;
+import com.google.cloud.spanner.SpannerOptions;
+import com.google.cloud.spanner.Statement;
+import com.google.spanner.v1.DirectedReadOptions;
+import com.google.spanner.v1.DirectedReadOptions.ExcludeReplicas;
+import com.google.spanner.v1.DirectedReadOptions.IncludeReplicas;
+import com.google.spanner.v1.DirectedReadOptions.ReplicaSelection;
+
+public class DirectedReadSample {
+  static void directedRead() {
+    // TODO(developer): Replace these variables before running the sample.
+    final String projectId = "my-project";
+    final String instanceId = "my-instance";
+    final String databaseId = "my-database";
+    directedRead(projectId, instanceId, databaseId);
+  }
+
+  static void directedRead(String projectId, String instanceId, String databaseId) {
+    // Only one of excludeReplicas or includeReplicas can be set
+    // Each accepts a list of replicaSelections which contains location and type
+    //   * `location` - The location must be one of the regions within the
+    //      multi-region configuration of your database.
+    //   * `type` - The type of the replica
+    // Some examples of using replicaSelectors are:
+    //   * `location:us-east1` --> The "us-east1" replica(s) of any available type
+    //                             will be used to process the request.
+    //   * `type:READ_ONLY`    --> The "READ_ONLY" type replica(s) in nearest
+    // .                            available location will be used to process the
+    //                             request.
+    //   * `location:us-east1 type:READ_ONLY` --> The "READ_ONLY" type replica(s)
+    //                          in location "us-east1" will be used to process
+    //                          the request.
+    //  includeReplicas also contains an option called autoFailoverDisabled, which when set to true
+    //  will instruct Spanner to not route requests to a replica outside the
+    //  includeReplicas list when all the specified replicas are unavailable
+    //  or unhealthy. Default value is `false`.
+    final DirectedReadOptions directedReadOptionsForClient =
+        DirectedReadOptions.newBuilder()
+            .setExcludeReplicas(
+                ExcludeReplicas.newBuilder()
+                    .addReplicaSelections(
+                        ReplicaSelection.newBuilder().setLocation("us-east4").build())
+                    .build())
+            .build();
+
+    // You can set default `DirectedReadOptions` for a Spanner client. These options will be applied
+    // to all read-only transactions that are executed by this client, unless specific
+    // DirectedReadOptions are set for a query.
+    // Directed read can only be used for read-only transactions. The default options will be
+    // ignored for any read/write transaction that the client executes.
+    try (Spanner spanner =
+        SpannerOptions.newBuilder()
+            .setProjectId(projectId)
+            .setDirectedReadOptions(directedReadOptionsForClient)
+            .build()
+            .getService()) {
+      final DatabaseClient dbClient =
+          spanner.getDatabaseClient(DatabaseId.of(projectId, instanceId, databaseId));
+
+      // DirectedReadOptions at request level will override the options set at
+      // client level (through SpannerOptions).
+      final DirectedReadOptions directedReadOptionsForRequest =
+          DirectedReadOptions.newBuilder()
+              .setIncludeReplicas(
+                  IncludeReplicas.newBuilder()
+                      .addReplicaSelections(
+                          ReplicaSelection.newBuilder()
+                              .setType(ReplicaSelection.Type.READ_WRITE)
+                              .build())
+                      .setAutoFailoverDisabled(true)
+                      .build())
+              .build();
+
+      // Read rows while passing DirectedReadOptions directly to the query.
+      try (ResultSet rs =
+          dbClient
+              .singleUse()
+              .executeQuery(
+                  Statement.of("SELECT SingerId, AlbumId, AlbumTitle FROM Albums"),
+                  Options.directedRead(directedReadOptionsForRequest))) {
+        while (rs.next()) {
+          System.out.printf(
+              "SingerId: %d, AlbumId: %d, AlbumTitle: %s\n",
+              rs.getLong(0), rs.getLong(1), rs.getString(2));
+        }
+        System.out.println("Successfully executed read-only transaction with directedReadOptions");
+      }
+    }
+  }
+}
+// [END spanner_directed_read]
diff --git a/samples/snippets/src/test/java/com/example/spanner/DirectedReadSampleIT.java b/samples/snippets/src/test/java/com/example/spanner/DirectedReadSampleIT.java
new file mode 100644
index 00000000000..771e157a3df
--- /dev/null
+++ b/samples/snippets/src/test/java/com/example/spanner/DirectedReadSampleIT.java
@@ -0,0 +1,105 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.example.spanner;
+
+import static com.example.spanner.SampleRunner.runSample;
+import static org.junit.Assert.assertTrue;
+
+import com.google.cloud.spanner.DatabaseAdminClient;
+import com.google.cloud.spanner.DatabaseClient;
+import com.google.cloud.spanner.DatabaseId;
+import com.google.cloud.spanner.KeySet;
+import com.google.cloud.spanner.Mutation;
+import com.google.cloud.spanner.Spanner;
+import com.google.cloud.spanner.SpannerOptions;
+import com.google.common.collect.ImmutableList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.concurrent.TimeUnit;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Integration tests for {@link DirectedReadSample} */
+@RunWith(JUnit4.class)
+public class DirectedReadSampleIT extends SampleTestBase {
+
+  private static DatabaseId databaseId;
+  private static Spanner spanner;
+
+  @BeforeClass
+  public static void createTestDatabase() throws Exception {
+    spanner = SpannerOptions.newBuilder().setProjectId(projectId).build().getService();
+    DatabaseAdminClient databaseAdminClient = spanner.getDatabaseAdminClient();
+    final String database = idGenerator.generateDatabaseId();
+    databaseAdminClient
+        .createDatabase(
+            instanceId,
+            database,
+            ImmutableList.of(
+                "CREATE TABLE Albums ("
+                    + "  SingerId    INT64 NOT NULL,"
+                    + "  AlbumId     INT64,"
+                    + "  AlbumTitle  STRING(1024)"
+                    + ") PRIMARY KEY (SingerId, AlbumId)"))
+        .get(10, TimeUnit.MINUTES);
+    databaseId = DatabaseId.of(projectId, instanceId, database);
+  }
+
+  @Before
+  public void insertTestData() {
+    final DatabaseClient client = spanner.getDatabaseClient(databaseId);
+    client.write(
+        Arrays.asList(
+            Mutation.newInsertOrUpdateBuilder("Albums")
+                .set("SingerId")
+                .to(1L)
+                .set("AlbumId")
+                .to(1L)
+                .set("AlbumTitle")
+                .to("title 1")
+                .build(),
+            Mutation.newInsertOrUpdateBuilder("Albums")
+                .set("SingerId")
+                .to(2L)
+                .set("AlbumId")
+                .to(2L)
+                .set("AlbumTitle")
+                .to("title 2")
+                .build()));
+  }
+
+  @After
+  public void removeTestData() {
+    final DatabaseClient client = spanner.getDatabaseClient(databaseId);
+    client.write(Collections.singletonList(Mutation.delete("Albums", KeySet.all())));
+  }
+
+  @Test
+  public void testDirectedRead() throws Exception {
+    final String out =
+        runSample(
+            () -> DirectedReadSample.directedRead(projectId, instanceId, databaseId.getDatabase()));
+    assertTrue(out.contains("SingerId: 1, AlbumId: 1, AlbumTitle: title 1"));
+    assertTrue(out.contains("SingerId: 2, AlbumId: 2, AlbumTitle: title 2"));
+    assertTrue(
+        out.contains("Successfully executed read-only transaction with directedReadOptions"));
+  }
+}