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 index 386eb780cd..4a57b76b1d 100644 --- 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 new file mode 100644 index 0000000000..d7a9d73d61 --- /dev/null +++ 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 index 4aed3439c9..ae4adc6de9 100644 --- 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)) {