Fix response codes returned by JSON formatting them (#2156) (#2200)

* Glue datasource support (#2055)



* Initial commit of new job APIs (#2050)



* Create Job API (#2070)

* Create Job API



* Refactor to Async Query API



---------



* Cancel Job API (#2126)



* Fix response codes returned



* Remove @opensearch datasource, update with new type



* Spotless Apply



* Fix tests



* Revert change back to json string



* Update tests to use JSON string literal instead of formatting



* Update IT



* Fix show datsources IT



* Remove files from merge



---------




(cherry picked from commit 5ce1bab)

Signed-off-by: Vamsi Manohar <reddyvam@amazon.com>
Signed-off-by: Derek Ho <dxho@amazon.com>
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>
Co-authored-by: Vamsi Manohar <reddyvam@amazon.com>

core/src/main/java/org/opensearch/sql/planner/physical/datasource/DataSourceTableScan.java

@@ -45,7 +45,7 @@ public String explain() {
   @Override
   public void open() {
     List<ExprValue> exprValues = new ArrayList<>();
-    Set<DataSourceMetadata> dataSourceMetadataSet = dataSourceService.getDataSourceMetadata(true);
+    Set<DataSourceMetadata> dataSourceMetadataSet = dataSourceService.getDataSourceMetadata(false);
     for (DataSourceMetadata dataSourceMetadata : dataSourceMetadataSet) {
       exprValues.add(
           new ExprTupleValue(

core/src/test/java/org/opensearch/sql/planner/physical/datasource/DataSourceTableScanTest.java

@@ -66,7 +66,7 @@ void testIterator() {
                         Collections.emptyList(),
                         ImmutableMap.of()))
             .collect(Collectors.toSet());
-    when(dataSourceService.getDataSourceMetadata(true)).thenReturn(dataSourceMetadata);
+    when(dataSourceService.getDataSourceMetadata(false)).thenReturn(dataSourceMetadata);
 
     assertFalse(dataSourceTableScan.hasNext());
     dataSourceTableScan.open();

datasources/src/main/java/org/opensearch/sql/datasources/transport/TransportCreateDataSourceAction.java

@@ -7,6 +7,8 @@
 
 package org.opensearch.sql.datasources.transport;
 
+import static org.opensearch.sql.protocol.response.format.JsonResponseFormatter.Style.PRETTY;
+
 import org.opensearch.action.ActionType;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.action.support.HandledTransportAction;
@@ -17,6 +19,7 @@
 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.protocol.response.format.JsonResponseFormatter;
 import org.opensearch.tasks.Task;
 import org.opensearch.transport.TransportService;
 
@@ -56,9 +59,14 @@ protected void doExecute(
     try {
       DataSourceMetadata dataSourceMetadata = request.getDataSourceMetadata();
       dataSourceService.createDataSource(dataSourceMetadata);
-      actionListener.onResponse(
-          new CreateDataSourceActionResponse(
-              "Created DataSource with name " + dataSourceMetadata.getName()));
+      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(responseContent));
     } catch (Exception e) {
       actionListener.onFailure(e);
     }

datasources/src/main/java/org/opensearch/sql/datasources/transport/TransportUpdateDataSourceAction.java

@@ -7,6 +7,8 @@
 
 package org.opensearch.sql.datasources.transport;
 
+import static org.opensearch.sql.protocol.response.format.JsonResponseFormatter.Style.PRETTY;
+
 import org.opensearch.action.ActionType;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.action.support.HandledTransportAction;
@@ -16,6 +18,7 @@
 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;
 
@@ -55,9 +58,14 @@ protected void doExecute(
       ActionListener<UpdateDataSourceActionResponse> actionListener) {
     try {
       dataSourceService.updateDataSource(request.getDataSourceMetadata());
-      actionListener.onResponse(
-          new UpdateDataSourceActionResponse(
-              "Updated DataSource with name " + request.getDataSourceMetadata().getName()));
+      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(responseContent));
     } catch (Exception e) {
       actionListener.onFailure(e);
     }

datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportCreateDataSourceActionTest.java

@@ -59,8 +59,8 @@ public void testDoExecute() {
         .onResponse(createDataSourceActionResponseArgumentCaptor.capture());
     CreateDataSourceActionResponse createDataSourceActionResponse =
         createDataSourceActionResponseArgumentCaptor.getValue();
-    Assertions.assertEquals(
-        "Created DataSource with name test_datasource", createDataSourceActionResponse.getResult());
+    String responseAsJson = "\"Created DataSource with name test_datasource\"";
+    Assertions.assertEquals(responseAsJson, createDataSourceActionResponse.getResult());
   }
 
   @Test

datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportUpdateDataSourceActionTest.java

@@ -59,8 +59,9 @@ public void testDoExecute() {
         .onResponse(updateDataSourceActionResponseArgumentCaptor.capture());
     UpdateDataSourceActionResponse updateDataSourceActionResponse =
         updateDataSourceActionResponseArgumentCaptor.getValue();
-    Assertions.assertEquals(
-        "Updated DataSource with name test_datasource", updateDataSourceActionResponse.getResult());
+    String responseAsJson = "\"Updated DataSource with name test_datasource\"";
+
+    Assertions.assertEquals(responseAsJson, updateDataSourceActionResponse.getResult());
   }
 
   @Test

integ-test/src/test/java/org/opensearch/sql/datasource/DataSourceAPIsIT.java

@@ -67,7 +67,7 @@ public void createDataSourceAPITest() {
     Response response = client().performRequest(createRequest);
     Assert.assertEquals(201, response.getStatusLine().getStatusCode());
     String createResponseString = getResponseBody(response);
-    Assert.assertEquals("Created DataSource with name create_prometheus", createResponseString);
+    Assert.assertEquals("\"Created DataSource with name create_prometheus\"", createResponseString);
     // Datasource is not immediately created. so introducing a sleep of 2s.
     Thread.sleep(2000);
 
@@ -109,7 +109,7 @@ public void updateDataSourceAPITest() {
     Response updateResponse = client().performRequest(updateRequest);
     Assert.assertEquals(200, updateResponse.getStatusLine().getStatusCode());
     String updateResponseString = getResponseBody(updateResponse);
-    Assert.assertEquals("Updated DataSource with name update_prometheus", updateResponseString);
+    Assert.assertEquals("\"Updated DataSource with name update_prometheus\"", updateResponseString);
 
     // Datasource is not immediately updated. so introducing a sleep of 2s.
     Thread.sleep(2000);

integ-test/src/test/java/org/opensearch/sql/ppl/ShowDataSourcesCommandIT.java

@@ -62,7 +62,7 @@ protected void deleteDataSourceMetadata() throws IOException {
   @Test
   public void testShowDataSourcesCommands() throws IOException {
     JSONObject result = executeQuery("show datasources");
-    verifyDataRows(result, rows("my_prometheus", "PROMETHEUS"), rows("@opensearch", "OPENSEARCH"));
+    verifyDataRows(result, rows("my_prometheus", "PROMETHEUS"));
     verifyColumn(result, columnName("DATASOURCE_NAME"), columnName("CONNECTOR_TYPE"));
   }