Initial try at using same permissions for patch and put

Signed-off-by: Derek Ho <[email protected]>

datasources/src/main/java/org/opensearch/sql/datasources/model/transport/DataSourceActionRequest.java

@@ -0,0 +1,37 @@
+package org.opensearch.sql.datasources.model.transport;
+
+import lombok.Getter;
+import org.opensearch.action.ActionRequest;
+import org.opensearch.action.ActionRequestValidationException;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.sql.datasource.model.DataSourceMetadata;
+
+import java.io.IOException;
+import java.util.Map;
+
+import static org.opensearch.sql.analysis.DataSourceSchemaIdentifierNameResolver.DEFAULT_DATASOURCE_NAME;
+import static org.opensearch.sql.datasources.utils.XContentParserUtils.CONNECTOR_FIELD;
+import static org.opensearch.sql.datasources.utils.XContentParserUtils.NAME_FIELD;
+
+public class DataSourceActionRequest extends ActionRequest {
+
+    public DataSourceActionRequest(StreamInput in) throws IOException {
+        super(in);
+    }
+
+    @Override
+    public ActionRequestValidationException validate() {
+        return null;
+    }
+
+    public DataSourceActionRequest() {
+    }
+
+//    public DataSourceMetadata getSourceMetadata() {
+//        return null;
+//    }
+//
+//    public Map<String, Object> getDataSourceData() {
+//        return Map.of();
+//    }
+}

datasources/src/main/java/org/opensearch/sql/datasources/model/transport/DataSourceActionResponse.java

@@ -0,0 +1,25 @@
+package org.opensearch.sql.datasources.model.transport;
+
+import lombok.Getter;
+import org.opensearch.core.action.ActionResponse;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+
+import java.io.IOException;
+
+public  class DataSourceActionResponse extends ActionResponse {
+    @Getter
+    String result;
+
+    public DataSourceActionResponse() {
+    }
+
+    public DataSourceActionResponse(StreamInput in) throws IOException {
+        super(in);
+    }
+
+    @Override
+    public void writeTo(StreamOutput streamOutput) throws IOException {
+
+    }
+}

datasources/src/main/java/org/opensearch/sql/datasources/model/transport/PatchDataSourceActionRequest.java

@@ -18,7 +18,7 @@
 import org.opensearch.action.ActionRequestValidationException;
 import org.opensearch.core.common.io.stream.StreamInput;
 
-public class PatchDataSourceActionRequest extends ActionRequest {
+public class PatchDataSourceActionRequest extends DataSourceActionRequest {
 
   @Getter private Map<String, Object> dataSourceData;
 

datasources/src/main/java/org/opensearch/sql/datasources/model/transport/PatchDataSourceActionResponse.java

@@ -15,7 +15,7 @@
 import org.opensearch.core.common.io.stream.StreamOutput;
 
 @RequiredArgsConstructor
-public class PatchDataSourceActionResponse extends ActionResponse {
+public class PatchDataSourceActionResponse extends DataSourceActionResponse {
 
   @Getter private final String result;
 

datasources/src/main/java/org/opensearch/sql/datasources/model/transport/UpdateDataSourceActionRequest.java

@@ -16,7 +16,7 @@
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.sql.datasource.model.DataSourceMetadata;
 
-public class UpdateDataSourceActionRequest extends ActionRequest {
+public class UpdateDataSourceActionRequest extends DataSourceActionRequest {
 
   @Getter private DataSourceMetadata dataSourceMetadata;
 

datasources/src/main/java/org/opensearch/sql/datasources/model/transport/UpdateDataSourceActionResponse.java

@@ -15,7 +15,7 @@
 import org.opensearch.core.common.io.stream.StreamOutput;
 
 @RequiredArgsConstructor
-public class UpdateDataSourceActionResponse extends ActionResponse {
+public class UpdateDataSourceActionResponse extends DataSourceActionResponse {
 
   @Getter private final String result;
 

datasources/src/main/java/org/opensearch/sql/datasources/rest/RestDataSourceQueryAction.java

@@ -197,20 +197,20 @@ private RestChannelConsumer executeUpdateRequest(RestRequest restRequest, NodeCl
             nodeClient,
             () ->
                 nodeClient.execute(
-                    TransportUpdateDataSourceAction.ACTION_TYPE,
+                    TransportDataSourceAction.ACTION_TYPE,
                     new UpdateDataSourceActionRequest(dataSourceMetadata),
                     new ActionListener<>() {
                       @Override
                       public void onResponse(
-                          UpdateDataSourceActionResponse updateDataSourceActionResponse) {
+                              DataSourceActionResponse updateDataSourceActionResponse) {
                         restChannel.sendResponse(
                             new BytesRestResponse(
                                 RestStatus.OK,
                                 "application/json; charset=UTF-8",
                                 updateDataSourceActionResponse.getResult()));
                       }
 
-                      @Override
+                        @Override
                       public void onFailure(Exception e) {
                         handleException(e, restChannel);
                       }
@@ -224,13 +224,14 @@ private RestChannelConsumer executePatchRequest(RestRequest restRequest, NodeCli
         Scheduler.schedule(
             nodeClient,
             () ->
+            {
                 nodeClient.execute(
-                    TransportPatchDataSourceAction.ACTION_TYPE,
+                        TransportDataSourceAction.ACTION_TYPE,
                     new PatchDataSourceActionRequest(dataSourceData),
                     new ActionListener<>() {
                       @Override
                       public void onResponse(
-                          PatchDataSourceActionResponse patchDataSourceActionResponse) {
+                          DataSourceActionResponse patchDataSourceActionResponse) {
                         restChannel.sendResponse(
                             new BytesRestResponse(
                                 RestStatus.OK,
@@ -242,7 +243,8 @@ public void onResponse(
                       public void onFailure(Exception e) {
                         handleException(e, restChannel);
                       }
-                    }));
+                    });
+            });
   }
 
   private RestChannelConsumer executeDeleteRequest(RestRequest restRequest, NodeClient nodeClient) {

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

@@ -0,0 +1,90 @@
+/*
+ *
+ *  * Copyright OpenSearch Contributors
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package org.opensearch.sql.datasources.transport;
+
+import static org.opensearch.sql.datasources.utils.XContentParserUtils.NAME_FIELD;
+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;
+import org.opensearch.common.inject.Inject;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.sql.datasource.DataSourceService;
+import org.opensearch.sql.datasources.model.transport.*;
+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;
+
+public class TransportDataSourceAction
+    extends HandledTransportAction<DataSourceActionRequest, DataSourceActionResponse> {
+
+  public static final String NAME = "cluster:admin/opensearch/ql/datasources/update";
+  public static final ActionType<DataSourceActionResponse> ACTION_TYPE =
+      new ActionType<>(NAME, DataSourceActionResponse::new);
+
+  private DataSourceService dataSourceService;
+
+  /**
+   * TransportUpdateDataSourceAction action for updating datasource.
+   *
+   * @param transportService transportService.
+   * @param actionFilters actionFilters.
+   * @param dataSourceService dataSourceService.
+   */
+  @Inject
+  public TransportDataSourceAction(
+      TransportService transportService,
+      ActionFilters actionFilters,
+      DataSourceServiceImpl dataSourceService) {
+    super(
+        TransportDataSourceAction.NAME,
+        transportService,
+        actionFilters,
+        DataSourceActionRequest::new);
+    this.dataSourceService = dataSourceService;
+  }
+
+  @Override
+  protected void doExecute(
+      Task task,
+      DataSourceActionRequest request,
+      ActionListener<DataSourceActionResponse> actionListener) {
+    try {
+      if (request instanceof UpdateDataSourceActionRequest) {
+        UpdateDataSourceActionRequest request1 = (UpdateDataSourceActionRequest) request;
+        dataSourceService.updateDataSource(request1.getDataSourceMetadata());
+        String responseContent =
+                new JsonResponseFormatter<String>(PRETTY) {
+                  @Override
+                  protected Object buildJsonObject(String response) {
+                    return response;
+                  }
+                }.format("Updated DataSource with name " + request1.getDataSourceMetadata().getName());
+        actionListener.onResponse(new UpdateDataSourceActionResponse(responseContent));
+      } else if (request instanceof PatchDataSourceActionRequest) {
+        PatchDataSourceActionRequest request2 = (PatchDataSourceActionRequest) request;
+        dataSourceService.patchDataSource(request2.getDataSourceData());
+        String responseContent =
+                new JsonResponseFormatter<String>(PRETTY) {
+                  @Override
+                  protected Object buildJsonObject(String response) {
+                    return response;
+                  }
+                }.format("Updated DataSource with name " + request2.getDataSourceData().get(NAME_FIELD));
+        actionListener.onResponse(new PatchDataSourceActionResponse(responseContent));
+      } else {
+        throw new IllegalArgumentException("Unexpected request type");
+      }
+//      actionListener.onResponse(new UpdateDataSourceActionResponse(responseContent));
+    } catch (Exception e) {
+      actionListener.onFailure(e);
+    }
+  }
+}

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

@@ -18,20 +18,21 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.sql.datasource.model.DataSourceMetadata;
 import org.opensearch.sql.datasource.model.DataSourceType;
+import org.opensearch.sql.datasources.model.transport.DataSourceActionResponse;
 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.tasks.Task;
 import org.opensearch.transport.TransportService;
 
 @ExtendWith(MockitoExtension.class)
-public class TransportUpdateDataSourceActionTest {
+public class TransportDataSourceActionTest {
 
   @Mock private TransportService transportService;
-  @Mock private TransportUpdateDataSourceAction action;
+  @Mock private TransportDataSourceAction action;
   @Mock private DataSourceServiceImpl dataSourceService;
   @Mock private Task task;
-  @Mock private ActionListener<UpdateDataSourceActionResponse> actionListener;
+  @Mock private ActionListener<DataSourceActionResponse> actionListener;
 
   @Captor
   private ArgumentCaptor<UpdateDataSourceActionResponse>
@@ -42,7 +43,7 @@ public class TransportUpdateDataSourceActionTest {
   @BeforeEach
   public void setUp() {
     action =
-        new TransportUpdateDataSourceAction(
+        new TransportDataSourceAction(
             transportService, new ActionFilters(new HashSet<>()), dataSourceService);
   }
 

plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java

@@ -170,12 +170,12 @@ public List<RestHandler> getRestHandlers(
             TransportGetDataSourceAction.class),
         new ActionHandler<>(
             new ActionType<>(
-                TransportUpdateDataSourceAction.NAME, UpdateDataSourceActionResponse::new),
-            TransportUpdateDataSourceAction.class),
-        new ActionHandler<>(
-            new ActionType<>(
-                TransportPatchDataSourceAction.NAME, PatchDataSourceActionResponse::new),
-            TransportPatchDataSourceAction.class),
+                TransportDataSourceAction.NAME, DataSourceActionResponse::new),
+            TransportDataSourceAction.class),
+//        new ActionHandler<>(
+//            new ActionType<>(
+//                TransportPatchDataSourceAction.NAME, PatchDataSourceActionResponse::new),
+//            TransportPatchDataSourceAction.class),
         new ActionHandler<>(
             new ActionType<>(
                 TransportDeleteDataSourceAction.NAME, DeleteDataSourceActionResponse::new),