From 056f54f3cc10d103151fccba569d46796a103591 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Wed, 6 Jan 2021 01:08:41 +0100
Subject: [PATCH] docs: add sample for timeout for one RPC (#707)

* docs: add sample for timeout for one RPC

Adds a sample for setting a timeout for a single RPC.

Fixes https://github.com/GoogleCloudPlatform/java-docs-samples/issues/3805

* fix: fix import order

* fix: format code with the correct formatter plugin

* fix: delete test data after each test

* fix: auto-throttle admin requests
---
 .../spanner/StatementTimeoutExample.java      | 82 +++++++++++++++++++
 .../spanner/SpannerStandaloneExamplesIT.java  | 10 +++
 2 files changed, 92 insertions(+)
 create mode 100644 samples/snippets/src/main/java/com/example/spanner/StatementTimeoutExample.java

diff --git a/samples/snippets/src/main/java/com/example/spanner/StatementTimeoutExample.java b/samples/snippets/src/main/java/com/example/spanner/StatementTimeoutExample.java
new file mode 100644
index 0000000000..9c3a8e5021
--- /dev/null
+++ b/samples/snippets/src/main/java/com/example/spanner/StatementTimeoutExample.java
@@ -0,0 +1,82 @@
+/*
+ * Copyright 2020 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;
+
+import com.google.api.gax.grpc.GrpcCallContext;
+import com.google.api.gax.rpc.ApiCallContext;
+import com.google.cloud.spanner.DatabaseClient;
+import com.google.cloud.spanner.DatabaseId;
+import com.google.cloud.spanner.Spanner;
+import com.google.cloud.spanner.SpannerOptions;
+import com.google.cloud.spanner.SpannerOptions.CallContextConfigurator;
+import com.google.cloud.spanner.Statement;
+import com.google.cloud.spanner.TransactionContext;
+import com.google.cloud.spanner.TransactionRunner.TransactionCallable;
+import com.google.spanner.v1.SpannerGrpc;
+import io.grpc.CallOptions;
+import io.grpc.Context;
+import io.grpc.MethodDescriptor;
+import java.util.concurrent.TimeUnit;
+
+class StatementTimeoutExample {
+
+  static void executeSqlWithTimeout() {
+    // TODO(developer): Replace these variables before running the sample.
+    String projectId = "my-project";
+    String instanceId = "my-instance";
+    String databaseId = "my-database";
+
+    try (Spanner spanner =
+        SpannerOptions.newBuilder().setProjectId(projectId).build().getService()) {
+      DatabaseClient client =
+          spanner.getDatabaseClient(DatabaseId.of(projectId, instanceId, databaseId));
+      executeSqlWithTimeout(client);
+    }
+  }
+
+  static void executeSqlWithTimeout(DatabaseClient client) {
+    CallContextConfigurator configurator = new CallContextConfigurator() {
+      public <ReqT, RespT> ApiCallContext configure(ApiCallContext context, ReqT request,
+          MethodDescriptor<ReqT, RespT> method) {
+        // DML uses the ExecuteSql RPC.
+        if (method == SpannerGrpc.getExecuteSqlMethod()) {
+          return GrpcCallContext.createDefault()
+              .withCallOptions(CallOptions.DEFAULT.withDeadlineAfter(60L, TimeUnit.SECONDS));
+        }
+        // Return null to indicate that the default should be used for other methods.
+        return null;
+      }
+    };
+    // Create a context that uses the custom call configuration.
+    Context context =
+        Context.current().withValue(SpannerOptions.CALL_CONTEXT_CONFIGURATOR_KEY, configurator);
+    // Run the transaction in the custom context.
+    context.run(new Runnable() {
+      public void run() {
+        client.readWriteTransaction().run(new TransactionCallable<long[]>() {
+          public long[] run(TransactionContext transaction) throws Exception {
+            String sql = "INSERT Singers (SingerId, FirstName, LastName)\n"
+                + "VALUES (20, 'George', 'Washington')";
+            long rowCount = transaction.executeUpdate(Statement.of(sql));
+            System.out.printf("%d record inserted.%n", rowCount);
+            return null;
+          }
+        });
+      }
+    });
+  }
+}
diff --git a/samples/snippets/src/test/java/com/example/spanner/SpannerStandaloneExamplesIT.java b/samples/snippets/src/test/java/com/example/spanner/SpannerStandaloneExamplesIT.java
index e8d1fdcb4c..2b3e016d1e 100644
--- a/samples/snippets/src/test/java/com/example/spanner/SpannerStandaloneExamplesIT.java
+++ b/samples/snippets/src/test/java/com/example/spanner/SpannerStandaloneExamplesIT.java
@@ -107,6 +107,7 @@ public void deleteTestData() {
     String projectId = spanner.getOptions().getProjectId();
     DatabaseClient client =
         spanner.getDatabaseClient(DatabaseId.of(projectId, instanceId, databaseId));
+    client.write(Collections.singleton(Mutation.delete("Singers", KeySet.all())));
     client.write(Collections.singleton(Mutation.delete("Venues", KeySet.all())));
   }
 
@@ -121,6 +122,15 @@ public void executeSqlWithCustomTimeoutAndRetrySettings_shouldWriteData() {
     assertThat(out).contains("1 record inserted.");
   }
 
+  @Test
+  public void executeSqlWithTimeout_shouldWriteData() {
+    String projectId = spanner.getOptions().getProjectId();
+    DatabaseClient client =
+        spanner.getDatabaseClient(DatabaseId.of(projectId, instanceId, databaseId));
+    String out = runExample(() -> StatementTimeoutExample.executeSqlWithTimeout(client));
+    assertThat(out).contains("1 record inserted.");
+  }
+
   @Test
   public void addNumericColumn_shouldSuccessfullyAddColumn()
       throws InterruptedException, ExecutionException {