From 80c251e6960491149d704310d2f83500243ec738 Mon Sep 17 00:00:00 2001
From: Vamsi Manohar <reddyvam@amazon.com>
Date: Sun, 18 Jun 2023 21:33:53 -0700
Subject: [PATCH] Improved error codes in case of security exception

Signed-off-by: Vamsi Manohar <reddyvam@amazon.com>
---
 .../opensearch/sql/datasources/exceptions/ErrorMessage.java   | 4 +++-
 .../sql/datasources/rest/RestDataSourceQueryAction.java       | 4 ++++
 2 files changed, 7 insertions(+), 1 deletion(-)

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 6dbd9bcfb5..c96640693c 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
@@ -7,6 +7,7 @@
 package org.opensearch.sql.datasources.exceptions;
 
 import com.google.gson.Gson;
+import com.google.gson.GsonBuilder;
 import com.google.gson.JsonObject;
 import lombok.Getter;
 import org.opensearch.rest.RestStatus;
@@ -65,7 +66,8 @@ public String toString() {
     JsonObject jsonObject = new JsonObject();
     jsonObject.addProperty("status", status);
     jsonObject.add("error", getErrorAsJson());
-    return new Gson().toJson(jsonObject);
+    Gson gson = new GsonBuilder().setPrettyPrinting().create();
+    return gson.toJson(jsonObject);
   }
 
   private JsonObject getErrorAsJson() {
diff --git a/datasources/src/main/java/org/opensearch/sql/datasources/rest/RestDataSourceQueryAction.java b/datasources/src/main/java/org/opensearch/sql/datasources/rest/RestDataSourceQueryAction.java
index c75170c355..95efd2e8f5 100644
--- a/datasources/src/main/java/org/opensearch/sql/datasources/rest/RestDataSourceQueryAction.java
+++ b/datasources/src/main/java/org/opensearch/sql/datasources/rest/RestDataSourceQueryAction.java
@@ -21,6 +21,7 @@
 import java.util.Locale;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.opensearch.OpenSearchException;
 import org.opensearch.action.ActionListener;
 import org.opensearch.client.node.NodeClient;
 import org.opensearch.rest.BaseRestHandler;
@@ -224,6 +225,9 @@ public void onFailure(Exception e) {
   private void handleException(Exception e, RestChannel restChannel) {
     if (e instanceof DataSourceNotFoundException) {
       reportError(restChannel, e, NOT_FOUND);
+    } else if (e instanceof OpenSearchException) {
+      OpenSearchException exception = (OpenSearchException) e;
+      reportError(restChannel, exception, exception.status());
     } else {
       LOG.error("Error happened during request handling", e);
       if (isClientError(e)) {