From ee6548cfa511d6efc99f508290ed0b1ce025a4cc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Fri, 2 Jun 2023 15:37:56 +0200
Subject: [PATCH] feat: make administrative request retries optional (#2476)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* feat: make administrative request retries optional

Make retries of administrative requests that fail because the
administrative request limit has been exceeded optional. This allows
users to apply custom retries or other custom handling for these errors.

* 🦉 Updates from OwlBot post-processor

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

---------

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
---
 .../google/cloud/spanner/SpannerOptions.java  | 15 ++++++++++-
 .../spanner/DatabaseAdminClientTest.java      | 27 +++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java
index 6663d3f8e6..33fa826bfe 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java
@@ -993,13 +993,26 @@ public Builder setAutoThrottleAdministrativeRequests() {
       return this;
     }
 
+    /**
+     * Disables automatic retries of administrative requests that fail if the <a
+     * href="https://cloud.google.com/spanner/quotas#administrative_limits">https://cloud.google.com/spanner/quotas#administrative_limits</a>
+     * have been exceeded. You should disable these retries if you intend to handle these errors in
+     * your application.
+     */
+    public Builder disableAdministrativeRequestRetries() {
+      this.retryAdministrativeRequestsSettings =
+          this.retryAdministrativeRequestsSettings.toBuilder().setMaxAttempts(1).build();
+      return this;
+    }
+
     /**
      * Sets the retry settings for retrying administrative requests when the quote of administrative
      * requests per minute has been exceeded.
      */
     Builder setRetryAdministrativeRequestsSettings(
         RetrySettings retryAdministrativeRequestsSettings) {
-      this.retryAdministrativeRequestsSettings = retryAdministrativeRequestsSettings;
+      this.retryAdministrativeRequestsSettings =
+          Preconditions.checkNotNull(retryAdministrativeRequestsSettings);
       return this;
     }
 
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseAdminClientTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseAdminClientTest.java
index eee9f520a7..5ef7b50804 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseAdminClientTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseAdminClientTest.java
@@ -935,4 +935,31 @@ public void testRetryOperationOnAdminMethodQuotaPerMinutePerProjectExceeded() {
     assertEquals(DB_ID, database.getId().getDatabase());
     assertEquals(2, mockDatabaseAdmin.countRequestsOfType(GetDatabaseRequest.class));
   }
+
+  @Test
+  public void testRetriesDisabledForOperationOnAdminMethodQuotaPerMinutePerProjectExceeded() {
+    ErrorInfo info =
+        ErrorInfo.newBuilder()
+            .putMetadata("quota_limit", "AdminMethodQuotaPerMinutePerProject")
+            .build();
+    Metadata.Key<ErrorInfo> key =
+        Metadata.Key.of(
+            info.getDescriptorForType().getFullName() + Metadata.BINARY_HEADER_SUFFIX,
+            ProtoLiteUtils.metadataMarshaller(info));
+    Metadata trailers = new Metadata();
+    trailers.put(key, info);
+    mockDatabaseAdmin.addException(
+        Status.RESOURCE_EXHAUSTED.withDescription("foo").asRuntimeException(trailers));
+    mockDatabaseAdmin.clearRequests();
+
+    Spanner spannerWithoutRetries =
+        spanner.getOptions().toBuilder().disableAdministrativeRequestRetries().build().getService();
+    AdminRequestsPerMinuteExceededException exception =
+        assertThrows(
+            AdminRequestsPerMinuteExceededException.class,
+            () -> spannerWithoutRetries.getDatabaseAdminClient().getDatabase(INSTANCE_ID, DB_ID));
+    assertEquals(ErrorCode.RESOURCE_EXHAUSTED, exception.getErrorCode());
+    // There should be only one request on the server, as the request was not retried.
+    assertEquals(1, mockDatabaseAdmin.countRequestsOfType(GetDatabaseRequest.class));
+  }
 }