From 6783ed7ac001c5b4e1c4da1ddab8b182e340e920 Mon Sep 17 00:00:00 2001
From: Derek Ho <dxho@amazon.com>
Date: Tue, 26 Sep 2023 14:12:43 -0400
Subject: [PATCH] Fix response codes returned

Signed-off-by: Derek Ho <dxho@amazon.com>
---
 .../transport/TransportCreateDataSourceAction.java | 14 ++++++++++++--
 .../transport/TransportUpdateDataSourceAction.java | 13 +++++++++++--
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/datasources/src/main/java/org/opensearch/sql/datasources/transport/TransportCreateDataSourceAction.java b/datasources/src/main/java/org/opensearch/sql/datasources/transport/TransportCreateDataSourceAction.java
index 54ca92b695..a05cbbac75 100644
--- a/datasources/src/main/java/org/opensearch/sql/datasources/transport/TransportCreateDataSourceAction.java
+++ b/datasources/src/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/datasources/src/main/java/org/opensearch/sql/datasources/transport/TransportUpdateDataSourceAction.java b/datasources/src/main/java/org/opensearch/sql/datasources/transport/TransportUpdateDataSourceAction.java
index 4325282f83..7e3807be62 100644
--- a/datasources/src/main/java/org/opensearch/sql/datasources/transport/TransportUpdateDataSourceAction.java
+++ b/datasources/src/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);
     }