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

* Glue datasource support (#2055) Signed-off-by: Vamsi Manohar <[email protected]> * Initial commit of new job APIs (#2050) Signed-off-by: Vamsi Manohar <[email protected]> * Create Job API (#2070) * Create Job API Signed-off-by: Vamsi Manohar <[email protected]> * Refactor to Async Query API Signed-off-by: Vamsi Manohar <[email protected]> --------- Signed-off-by: Vamsi Manohar <[email protected]> * Cancel Job API (#2126) Signed-off-by: Vamsi Manohar <[email protected]> * Fix response codes returned Signed-off-by: Derek Ho <[email protected]> * Remove @opensearch datasource, update with new type Signed-off-by: Derek Ho <[email protected]> * Spotless Apply Signed-off-by: Derek Ho <[email protected]> * Fix tests Signed-off-by: Derek Ho <[email protected]> * Revert change back to json string Signed-off-by: Derek Ho <[email protected]> * Update tests to use JSON string literal instead of formatting Signed-off-by: Derek Ho <[email protected]> * Update IT Signed-off-by: Derek Ho <[email protected]> * Fix show datsources IT Signed-off-by: Derek Ho <[email protected]> * Remove files from merge Signed-off-by: Derek Ho <[email protected]> --------- Signed-off-by: Vamsi Manohar <[email protected]> Signed-off-by: Derek Ho <[email protected]> Co-authored-by: Vamsi Manohar <[email protected]>
opensearch-project · Oct 3, 2023 · 5ce1bab · 5ce1bab
1 parent eecff50
commit 5ce1bab
Show file tree

Hide file tree

Showing 8 changed files with 32 additions and 15 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/datasource/DataSourceTableScan.java b/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(

diff --git a/...src/test/java/org/opensearch/sql/planner/physical/datasource/DataSourceTableScanTest.java b/...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();

diff --git a/...c/main/java/org/opensearch/sql/datasources/transport/TransportCreateDataSourceAction.java b/...c/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);
     }

diff --git a/...c/main/java/org/opensearch/sql/datasources/transport/TransportUpdateDataSourceAction.java b/...c/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);
     }

diff --git a/...st/java/org/opensearch/sql/datasources/transport/TransportCreateDataSourceActionTest.java b/...st/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

diff --git a/...st/java/org/opensearch/sql/datasources/transport/TransportUpdateDataSourceActionTest.java b/...st/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

diff --git a/integ-test/src/test/java/org/opensearch/sql/datasource/DataSourceAPIsIT.java b/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);

diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/ShowDataSourcesCommandIT.java b/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"));
   }