Return 429 for ConcurrencyLimitExceededException (#2428) (#2429)

* Return 429 for ConcurrencyLimitExceededException * add UT --------- (cherry picked from commit 2c230be) Signed-off-by: Peng Huo <[email protected]> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
opensearch-project · Nov 30, 2023 · d1ed772 · d1ed772
1 parent 4af150e
commit d1ed772
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 0 deletions.
diff --git a/datasources/src/main/java/org/opensearch/sql/datasources/exceptions/ErrorMessage.java b/datasources/src/main/java/org/opensearch/sql/datasources/exceptions/ErrorMessage.java
@@ -39,6 +39,9 @@ private String fetchType() {
   }
 
   protected String fetchReason() {
+    if (status == RestStatus.TOO_MANY_REQUESTS.getStatus()) {
+      return "Too Many Requests";
+    }
     return status == RestStatus.BAD_REQUEST.getStatus()
         ? "Invalid Request"
         : "There was internal problem at backend";

diff --git a/datasources/src/test/java/org/opensearch/sql/datasources/exceptions/ErrorMessageTest.java b/datasources/src/test/java/org/opensearch/sql/datasources/exceptions/ErrorMessageTest.java
@@ -0,0 +1,21 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.datasources.exceptions;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import org.junit.jupiter.api.Test;
+import org.opensearch.core.rest.RestStatus;
+
+class ErrorMessageTest {
+
+  @Test
+  void fetchReason() {
+    ErrorMessage errorMessage =
+        new ErrorMessage(new RuntimeException(), RestStatus.TOO_MANY_REQUESTS.getStatus());
+    assertEquals("Too Many Requests", errorMessage.getReason());
+  }
+}
diff --git a/spark/src/main/java/org/opensearch/sql/spark/rest/RestAsyncQueryManagementAction.java b/spark/src/main/java/org/opensearch/sql/spark/rest/RestAsyncQueryManagementAction.java
@@ -7,6 +7,7 @@
 
 import static org.opensearch.core.rest.RestStatus.BAD_REQUEST;
 import static org.opensearch.core.rest.RestStatus.SERVICE_UNAVAILABLE;
+import static org.opensearch.core.rest.RestStatus.TOO_MANY_REQUESTS;
 import static org.opensearch.rest.RestRequest.Method.DELETE;
 import static org.opensearch.rest.RestRequest.Method.GET;
 import static org.opensearch.rest.RestRequest.Method.POST;
@@ -29,6 +30,7 @@
 import org.opensearch.sql.datasources.utils.Scheduler;
 import org.opensearch.sql.legacy.metrics.MetricName;
 import org.opensearch.sql.legacy.utils.MetricUtils;
+import org.opensearch.sql.spark.leasemanager.ConcurrencyLimitExceededException;
 import org.opensearch.sql.spark.rest.model.CreateAsyncQueryRequest;
 import org.opensearch.sql.spark.transport.TransportCancelAsyncQueryRequestAction;
 import org.opensearch.sql.spark.transport.TransportCreateAsyncQueryRequestAction;
@@ -175,6 +177,10 @@ private void handleException(
       OpenSearchException exception = (OpenSearchException) e;
       reportError(restChannel, exception, exception.status());
       addCustomerErrorMetric(requestMethod);
+    } else if (e instanceof ConcurrencyLimitExceededException) {
+      LOG.error("Too many request", e);
+      reportError(restChannel, e, TOO_MANY_REQUESTS);
+      addCustomerErrorMetric(requestMethod);
     } else {
       LOG.error("Error happened during request handling", e);
       if (isClientError(e)) {