Fix response codes returned

Signed-off-by: Derek Ho <[email protected]>
opensearch-project · Sep 26, 2023 · 6783ed7 · 6783ed7
1 parent be82714
commit 6783ed7
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 4 deletions.
diff --git a/...c/main/java/org/opensearch/sql/datasources/transport/TransportCreateDataSourceAction.java b/...c/main/java/org/opensearch/sql/datasources/transport/TransportCreateDataSourceAction.java
@@ -17,9 +17,13 @@
 import org.opensearch.sql.datasources.model.transport.CreateDataSourceActionRequest;
 import org.opensearch.sql.datasources.model.transport.CreateDataSourceActionResponse;
 import org.opensearch.sql.datasources.service.DataSourceServiceImpl;
+import org.opensearch.sql.executor.ExecutionEngine;
+import org.opensearch.sql.protocol.response.format.JsonResponseFormatter;
 import org.opensearch.tasks.Task;
 import org.opensearch.transport.TransportService;
 
+import static org.opensearch.sql.protocol.response.format.JsonResponseFormatter.Style.PRETTY;
+
 public class TransportCreateDataSourceAction
     extends HandledTransportAction<CreateDataSourceActionRequest, CreateDataSourceActionResponse> {
   public static final String NAME = "cluster:admin/opensearch/ql/datasources/create";
@@ -56,9 +60,15 @@ protected void doExecute(
     try {
       DataSourceMetadata dataSourceMetadata = request.getDataSourceMetadata();
       dataSourceService.createDataSource(dataSourceMetadata);
+      String responseContent =
+              new JsonResponseFormatter<String>(PRETTY) {
+                @Override
+                protected Object buildJsonObject(String response) {
+                  return response;
+                }
+              }.format("Created DataSource with name " + dataSourceMetadata.getName());
       actionListener.onResponse(
-          new CreateDataSourceActionResponse(
-              "Created DataSource with name " + dataSourceMetadata.getName()));
+          new CreateDataSourceActionResponse(responseContent));
     } catch (Exception e) {
       actionListener.onFailure(e);
     }

diff --git a/...c/main/java/org/opensearch/sql/datasources/transport/TransportUpdateDataSourceAction.java b/...c/main/java/org/opensearch/sql/datasources/transport/TransportUpdateDataSourceAction.java
@@ -16,9 +16,12 @@
 import org.opensearch.sql.datasources.model.transport.UpdateDataSourceActionRequest;
 import org.opensearch.sql.datasources.model.transport.UpdateDataSourceActionResponse;
 import org.opensearch.sql.datasources.service.DataSourceServiceImpl;
+import org.opensearch.sql.protocol.response.format.JsonResponseFormatter;
 import org.opensearch.tasks.Task;
 import org.opensearch.transport.TransportService;
 
+import static org.opensearch.sql.protocol.response.format.JsonResponseFormatter.Style.PRETTY;
+
 public class TransportUpdateDataSourceAction
     extends HandledTransportAction<UpdateDataSourceActionRequest, UpdateDataSourceActionResponse> {
 
@@ -55,9 +58,15 @@ protected void doExecute(
       ActionListener<UpdateDataSourceActionResponse> actionListener) {
     try {
       dataSourceService.updateDataSource(request.getDataSourceMetadata());
+      String responseContent =
+              new JsonResponseFormatter<String>(PRETTY) {
+                @Override
+                protected Object buildJsonObject(String response) {
+                  return response;
+                }
+              }.format("Updated DataSource with name " +  request.getDataSourceMetadata().getName());
       actionListener.onResponse(
-          new UpdateDataSourceActionResponse(
-              "Updated DataSource with name " + request.getDataSourceMetadata().getName()));
+          new UpdateDataSourceActionResponse(responseContent));
     } catch (Exception e) {
       actionListener.onFailure(e);
     }