diff --git a/core/src/main/java/org/opensearch/sql/datasource/DataSourceService.java b/core/src/main/java/org/opensearch/sql/datasource/DataSourceService.java index 6dace50f99..162fe9e8f8 100644 --- a/core/src/main/java/org/opensearch/sql/datasource/DataSourceService.java +++ b/core/src/main/java/org/opensearch/sql/datasource/DataSourceService.java @@ -5,6 +5,7 @@ package org.opensearch.sql.datasource; +import java.util.Map; import java.util.Set; import org.opensearch.sql.datasource.model.DataSource; import org.opensearch.sql.datasource.model.DataSourceMetadata; @@ -56,12 +57,19 @@ public interface DataSourceService { void createDataSource(DataSourceMetadata metadata); /** - * Updates {@link DataSource} corresponding to dataSourceMetadata. + * Updates {@link DataSource} corresponding to dataSourceMetadata (all fields needed). * * @param dataSourceMetadata {@link DataSourceMetadata}. */ void updateDataSource(DataSourceMetadata dataSourceMetadata); + /** + * Patches {@link DataSource} corresponding to the given name (only fields to be changed needed). + * + * @param dataSourceData + */ + void patchDataSource(Map dataSourceData); + /** * Deletes {@link DataSource} corresponding to the DataSource name. * diff --git a/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTestBase.java b/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTestBase.java index 508567582b..569cdd96f8 100644 --- a/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTestBase.java +++ b/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTestBase.java @@ -231,6 +231,9 @@ public DataSource getDataSource(String dataSourceName) { @Override public void updateDataSource(DataSourceMetadata dataSourceMetadata) {} + @Override + public void patchDataSource(Map dataSourceData) {} + @Override public void deleteDataSource(String dataSourceName) {} diff --git a/datasources/src/main/java/org/opensearch/sql/datasources/model/transport/PatchDataSourceActionRequest.java b/datasources/src/main/java/org/opensearch/sql/datasources/model/transport/PatchDataSourceActionRequest.java new file mode 100644 index 0000000000..9443ea561e --- /dev/null +++ b/datasources/src/main/java/org/opensearch/sql/datasources/model/transport/PatchDataSourceActionRequest.java @@ -0,0 +1,49 @@ +/* + * + * * Copyright OpenSearch Contributors + * * SPDX-License-Identifier: Apache-2.0 + * + */ + +package org.opensearch.sql.datasources.model.transport; + +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; + +import java.io.IOException; +import java.util.Map; +import lombok.Getter; +import org.opensearch.action.ActionRequest; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.core.common.io.stream.StreamInput; + +public class PatchDataSourceActionRequest extends ActionRequest { + + @Getter private Map dataSourceData; + + /** Constructor of UpdateDataSourceActionRequest from StreamInput. */ + public PatchDataSourceActionRequest(StreamInput in) throws IOException { + super(in); + } + + public PatchDataSourceActionRequest(Map dataSourceData) { + this.dataSourceData = dataSourceData; + } + + @Override + public ActionRequestValidationException validate() { + if (this.dataSourceData.get(NAME_FIELD).equals(DEFAULT_DATASOURCE_NAME)) { + ActionRequestValidationException exception = new ActionRequestValidationException(); + exception.addValidationError( + "Not allowed to update datasource with name : " + DEFAULT_DATASOURCE_NAME); + return exception; + } else if (this.dataSourceData.get(CONNECTOR_FIELD) != null) { + ActionRequestValidationException exception = new ActionRequestValidationException(); + exception.addValidationError("Not allowed to update connector for datasource"); + return exception; + } else { + return null; + } + } +} diff --git a/datasources/src/main/java/org/opensearch/sql/datasources/model/transport/PatchDataSourceActionResponse.java b/datasources/src/main/java/org/opensearch/sql/datasources/model/transport/PatchDataSourceActionResponse.java new file mode 100644 index 0000000000..18873a6731 --- /dev/null +++ b/datasources/src/main/java/org/opensearch/sql/datasources/model/transport/PatchDataSourceActionResponse.java @@ -0,0 +1,31 @@ +/* + * + * * Copyright OpenSearch Contributors + * * SPDX-License-Identifier: Apache-2.0 + * + */ + +package org.opensearch.sql.datasources.model.transport; + +import java.io.IOException; +import lombok.Getter; +import lombok.RequiredArgsConstructor; +import org.opensearch.core.action.ActionResponse; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; + +@RequiredArgsConstructor +public class PatchDataSourceActionResponse extends ActionResponse { + + @Getter private final String result; + + public PatchDataSourceActionResponse(StreamInput in) throws IOException { + super(in); + result = in.readString(); + } + + @Override + public void writeTo(StreamOutput streamOutput) throws IOException { + streamOutput.writeString(result); + } +} diff --git a/datasources/src/main/java/org/opensearch/sql/datasources/rest/RestDataSourceQueryAction.java b/datasources/src/main/java/org/opensearch/sql/datasources/rest/RestDataSourceQueryAction.java index 2947afc5b9..c207f55738 100644 --- a/datasources/src/main/java/org/opensearch/sql/datasources/rest/RestDataSourceQueryAction.java +++ b/datasources/src/main/java/org/opensearch/sql/datasources/rest/RestDataSourceQueryAction.java @@ -10,15 +10,13 @@ import static org.opensearch.core.rest.RestStatus.BAD_REQUEST; import static org.opensearch.core.rest.RestStatus.NOT_FOUND; import static org.opensearch.core.rest.RestStatus.SERVICE_UNAVAILABLE; -import static org.opensearch.rest.RestRequest.Method.DELETE; -import static org.opensearch.rest.RestRequest.Method.GET; -import static org.opensearch.rest.RestRequest.Method.POST; -import static org.opensearch.rest.RestRequest.Method.PUT; +import static org.opensearch.rest.RestRequest.Method.*; import com.google.common.collect.ImmutableList; import java.io.IOException; import java.util.List; import java.util.Locale; +import java.util.Map; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.opensearch.OpenSearchException; @@ -32,18 +30,8 @@ import org.opensearch.sql.datasource.model.DataSourceMetadata; import org.opensearch.sql.datasources.exceptions.DataSourceNotFoundException; import org.opensearch.sql.datasources.exceptions.ErrorMessage; -import org.opensearch.sql.datasources.model.transport.CreateDataSourceActionRequest; -import org.opensearch.sql.datasources.model.transport.CreateDataSourceActionResponse; -import org.opensearch.sql.datasources.model.transport.DeleteDataSourceActionRequest; -import org.opensearch.sql.datasources.model.transport.DeleteDataSourceActionResponse; -import org.opensearch.sql.datasources.model.transport.GetDataSourceActionRequest; -import org.opensearch.sql.datasources.model.transport.GetDataSourceActionResponse; -import org.opensearch.sql.datasources.model.transport.UpdateDataSourceActionRequest; -import org.opensearch.sql.datasources.model.transport.UpdateDataSourceActionResponse; -import org.opensearch.sql.datasources.transport.TransportCreateDataSourceAction; -import org.opensearch.sql.datasources.transport.TransportDeleteDataSourceAction; -import org.opensearch.sql.datasources.transport.TransportGetDataSourceAction; -import org.opensearch.sql.datasources.transport.TransportUpdateDataSourceAction; +import org.opensearch.sql.datasources.model.transport.*; +import org.opensearch.sql.datasources.transport.*; import org.opensearch.sql.datasources.utils.Scheduler; import org.opensearch.sql.datasources.utils.XContentParserUtils; @@ -98,6 +86,17 @@ public List routes() { */ new Route(PUT, BASE_DATASOURCE_ACTION_URL), + /* + * PATCH datasources + * Request body: + * Ref + * [org.opensearch.sql.plugin.transport.datasource.model.PatchDataSourceActionRequest] + * Response body: + * Ref + * [org.opensearch.sql.plugin.transport.datasource.model.PatchDataSourceActionResponse] + */ + new Route(PATCH, BASE_DATASOURCE_ACTION_URL), + /* * DELETE datasources * Request body: Ref @@ -122,6 +121,8 @@ protected RestChannelConsumer prepareRequest(RestRequest restRequest, NodeClient return executeUpdateRequest(restRequest, nodeClient); case DELETE: return executeDeleteRequest(restRequest, nodeClient); + case PATCH: + return executePatchRequest(restRequest, nodeClient); default: return restChannel -> restChannel.sendResponse( @@ -216,6 +217,34 @@ public void onFailure(Exception e) { })); } + private RestChannelConsumer executePatchRequest(RestRequest restRequest, NodeClient nodeClient) + throws IOException { + Map dataSourceData = XContentParserUtils.toMap(restRequest.contentParser()); + return restChannel -> + Scheduler.schedule( + nodeClient, + () -> + nodeClient.execute( + TransportPatchDataSourceAction.ACTION_TYPE, + new PatchDataSourceActionRequest(dataSourceData), + new ActionListener<>() { + @Override + public void onResponse( + PatchDataSourceActionResponse patchDataSourceActionResponse) { + restChannel.sendResponse( + new BytesRestResponse( + RestStatus.OK, + "application/json; charset=UTF-8", + patchDataSourceActionResponse.getResult())); + } + + @Override + public void onFailure(Exception e) { + handleException(e, restChannel); + } + })); + } + private RestChannelConsumer executeDeleteRequest(RestRequest restRequest, NodeClient nodeClient) { String dataSourceName = restRequest.param("dataSourceName"); diff --git a/datasources/src/main/java/org/opensearch/sql/datasources/service/DataSourceServiceImpl.java b/datasources/src/main/java/org/opensearch/sql/datasources/service/DataSourceServiceImpl.java index 25e8006d66..8ba618fb44 100644 --- a/datasources/src/main/java/org/opensearch/sql/datasources/service/DataSourceServiceImpl.java +++ b/datasources/src/main/java/org/opensearch/sql/datasources/service/DataSourceServiceImpl.java @@ -6,15 +6,11 @@ package org.opensearch.sql.datasources.service; import static org.opensearch.sql.analysis.DataSourceSchemaIdentifierNameResolver.DEFAULT_DATASOURCE_NAME; +import static org.opensearch.sql.datasources.utils.XContentParserUtils.*; import com.google.common.base.Preconditions; import com.google.common.base.Strings; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Objects; -import java.util.Optional; -import java.util.Set; +import java.util.*; import org.opensearch.sql.common.utils.StringUtils; import org.opensearch.sql.datasource.DataSourceService; import org.opensearch.sql.datasource.model.DataSource; @@ -100,6 +96,19 @@ public void updateDataSource(DataSourceMetadata dataSourceMetadata) { } } + @Override + public void patchDataSource(Map dataSourceData) { + if (!dataSourceData.get(NAME_FIELD).equals(DEFAULT_DATASOURCE_NAME)) { + DataSourceMetadata dataSourceMetadata = + getRawDataSourceMetadata((String) dataSourceData.get(NAME_FIELD)); + replaceOldDatasourceMetadata(dataSourceData, dataSourceMetadata); + updateDataSource(dataSourceMetadata); + } else { + throw new UnsupportedOperationException( + "Not allowed to update default datasource :" + DEFAULT_DATASOURCE_NAME); + } + } + @Override public void deleteDataSource(String dataSourceName) { if (dataSourceName.equals(DEFAULT_DATASOURCE_NAME)) { @@ -136,6 +145,35 @@ private void validateDataSourceMetaData(DataSourceMetadata metadata) { + " Properties are required parameters."); } + /** + * Replaces the fields in the map of the given metadata. + * + * @param dataSourceData + * @param metadata {@link DataSourceMetadata}. + */ + private void replaceOldDatasourceMetadata( + Map dataSourceData, DataSourceMetadata metadata) { + + for (String key : dataSourceData.keySet()) { + switch (key) { + // Name and connector should not be modified + case DESCRIPTION_FIELD: + metadata.setDescription((String) dataSourceData.get(DESCRIPTION_FIELD)); + break; + case ALLOWED_ROLES_FIELD: + metadata.setAllowedRoles((List) dataSourceData.get(ALLOWED_ROLES_FIELD)); + break; + case PROPERTIES_FIELD: + Map properties = new HashMap<>(metadata.getProperties()); + properties.putAll(((Map) dataSourceData.get(PROPERTIES_FIELD))); + break; + case NAME_FIELD: + case CONNECTOR_FIELD: + break; + } + } + } + @Override public DataSourceMetadata getRawDataSourceMetadata(String dataSourceName) { if (dataSourceName.equals(DEFAULT_DATASOURCE_NAME)) { diff --git a/datasources/src/main/java/org/opensearch/sql/datasources/transport/TransportPatchDataSourceAction.java b/datasources/src/main/java/org/opensearch/sql/datasources/transport/TransportPatchDataSourceAction.java new file mode 100644 index 0000000000..303e905cec --- /dev/null +++ b/datasources/src/main/java/org/opensearch/sql/datasources/transport/TransportPatchDataSourceAction.java @@ -0,0 +1,74 @@ +/* + * + * * 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.PatchDataSourceActionRequest; +import org.opensearch.sql.datasources.model.transport.PatchDataSourceActionResponse; +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 TransportPatchDataSourceAction + extends HandledTransportAction { + + public static final String NAME = "cluster:admin/opensearch/ql/datasources/patch"; + public static final ActionType ACTION_TYPE = + new ActionType<>(NAME, PatchDataSourceActionResponse::new); + + private DataSourceService dataSourceService; + + /** + * TransportPatchDataSourceAction action for updating datasource. + * + * @param transportService transportService. + * @param actionFilters actionFilters. + * @param dataSourceService dataSourceService. + */ + @Inject + public TransportPatchDataSourceAction( + TransportService transportService, + ActionFilters actionFilters, + DataSourceServiceImpl dataSourceService) { + super( + TransportPatchDataSourceAction.NAME, + transportService, + actionFilters, + PatchDataSourceActionRequest::new); + this.dataSourceService = dataSourceService; + } + + @Override + protected void doExecute( + Task task, + PatchDataSourceActionRequest request, + ActionListener actionListener) { + try { + dataSourceService.patchDataSource(request.getDataSourceData()); + String responseContent = + new JsonResponseFormatter(PRETTY) { + @Override + protected Object buildJsonObject(String response) { + return response; + } + }.format("Updated DataSource with name " + request.getDataSourceData().get(NAME_FIELD)); + actionListener.onResponse(new PatchDataSourceActionResponse(responseContent)); + } catch (Exception e) { + actionListener.onFailure(e); + } + } +} diff --git a/datasources/src/main/java/org/opensearch/sql/datasources/utils/XContentParserUtils.java b/datasources/src/main/java/org/opensearch/sql/datasources/utils/XContentParserUtils.java index 261f13870a..6af2a5a761 100644 --- a/datasources/src/main/java/org/opensearch/sql/datasources/utils/XContentParserUtils.java +++ b/datasources/src/main/java/org/opensearch/sql/datasources/utils/XContentParserUtils.java @@ -90,6 +90,59 @@ public static DataSourceMetadata toDataSourceMetadata(XContentParser parser) thr name, description, connector, allowedRoles, properties, resultIndex); } + public static Map toMap(XContentParser parser) throws IOException { + Map resultMap = new HashMap<>(); + String name; + String description; + List allowedRoles = new ArrayList<>(); + Map properties = new HashMap<>(); + String resultIndex; + ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser); + while (parser.nextToken() != XContentParser.Token.END_OBJECT) { + String fieldName = parser.currentName(); + parser.nextToken(); + switch (fieldName) { + case NAME_FIELD: + name = parser.textOrNull(); + resultMap.put(NAME_FIELD, name); + break; + case DESCRIPTION_FIELD: + description = parser.textOrNull(); + resultMap.put(DESCRIPTION_FIELD, description); + break; + case CONNECTOR_FIELD: + // no-op - datasource connector should not be modified + break; + case ALLOWED_ROLES_FIELD: + while (parser.nextToken() != XContentParser.Token.END_ARRAY) { + allowedRoles.add(parser.text()); + } + resultMap.put(ALLOWED_ROLES_FIELD, allowedRoles); + break; + case PROPERTIES_FIELD: + ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser); + while (parser.nextToken() != XContentParser.Token.END_OBJECT) { + String key = parser.currentName(); + parser.nextToken(); + String value = parser.textOrNull(); + properties.put(key, value); + } + resultMap.put(PROPERTIES_FIELD, properties); + break; + case RESULT_INDEX_FIELD: + resultIndex = parser.textOrNull(); + resultMap.put(RESULT_INDEX_FIELD, resultIndex); + break; + default: + throw new IllegalArgumentException("Unknown field: " + fieldName); + } + } + if (resultMap.get(NAME_FIELD) == null || resultMap.get(NAME_FIELD) == "") { + throw new IllegalArgumentException("Name is a required field."); + } + return resultMap; + } + /** * Converts json string to DataSourceMetadata. * @@ -109,6 +162,25 @@ public static DataSourceMetadata toDataSourceMetadata(String json) throws IOExce } } + /** + * Converts json string to Map. + * + * @param json jsonstring. + * @return DataSourceData + * @throws IOException IOException. + */ + public static Map toMap(String json) throws IOException { + try (XContentParser parser = + XContentType.JSON + .xContent() + .createParser( + NamedXContentRegistry.EMPTY, + DeprecationHandler.THROW_UNSUPPORTED_OPERATION, + json)) { + return toMap(parser); + } + } + /** * Converts DataSourceMetadata to XContentBuilder. * diff --git a/datasources/src/test/java/org/opensearch/sql/datasources/service/DataSourceServiceImplTest.java b/datasources/src/test/java/org/opensearch/sql/datasources/service/DataSourceServiceImplTest.java index 6164d8b73f..c62e586dae 100644 --- a/datasources/src/test/java/org/opensearch/sql/datasources/service/DataSourceServiceImplTest.java +++ b/datasources/src/test/java/org/opensearch/sql/datasources/service/DataSourceServiceImplTest.java @@ -18,6 +18,7 @@ import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; import static org.opensearch.sql.analysis.DataSourceSchemaIdentifierNameResolver.DEFAULT_DATASOURCE_NAME; +import static org.opensearch.sql.datasources.utils.XContentParserUtils.*; import com.google.common.collect.ImmutableMap; import java.util.ArrayList; @@ -264,7 +265,6 @@ void testGetDataSourceMetadataSetWithDefaultDatasource() { @Test void testUpdateDataSourceSuccessCase() { - DataSourceMetadata dataSourceMetadata = metadata("testDS", DataSourceType.OPENSEARCH, Collections.emptyList(), ImmutableMap.of()); dataSourceService.updateDataSource(dataSourceMetadata); @@ -289,6 +289,46 @@ void testUpdateDefaultDataSource() { unsupportedOperationException.getMessage()); } + @Test + void testPatchDefaultDataSource() { + Map dataSourceData = + Map.of(NAME_FIELD, DEFAULT_DATASOURCE_NAME, DESCRIPTION_FIELD, "test"); + UnsupportedOperationException unsupportedOperationException = + assertThrows( + UnsupportedOperationException.class, + () -> dataSourceService.patchDataSource(dataSourceData)); + assertEquals( + "Not allowed to update default datasource :" + DEFAULT_DATASOURCE_NAME, + unsupportedOperationException.getMessage()); + } + + @Test + void testPatchDataSourceSuccessCase() { + // Tests that patch underlying implementation is to call update + Map dataSourceData = + new HashMap<>( + Map.of( + NAME_FIELD, + "testDS", + DESCRIPTION_FIELD, + "test", + CONNECTOR_FIELD, + "PROMETHEUS", + ALLOWED_ROLES_FIELD, + new ArrayList<>(), + PROPERTIES_FIELD, + Map.of(), + RESULT_INDEX_FIELD, + "")); + DataSourceMetadata getData = + metadata("testDS", DataSourceType.OPENSEARCH, Collections.emptyList(), ImmutableMap.of()); + when(dataSourceMetadataStorage.getDataSourceMetadata("testDS")) + .thenReturn(Optional.ofNullable(getData)); + + dataSourceService.patchDataSource(dataSourceData); + verify(dataSourceMetadataStorage, times(1)).updateDataSourceMetadata(any()); + } + @Test void testDeleteDatasource() { dataSourceService.deleteDataSource("testDS"); diff --git a/datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportPatchDataSourceActionTest.java b/datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportPatchDataSourceActionTest.java new file mode 100644 index 0000000000..5e1e7df112 --- /dev/null +++ b/datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportPatchDataSourceActionTest.java @@ -0,0 +1,79 @@ +package org.opensearch.sql.datasources.transport; + +import static org.mockito.Mockito.*; +import static org.opensearch.sql.datasources.utils.XContentParserUtils.*; + +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Captor; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.core.action.ActionListener; +import org.opensearch.sql.datasources.model.transport.PatchDataSourceActionRequest; +import org.opensearch.sql.datasources.model.transport.PatchDataSourceActionResponse; +import org.opensearch.sql.datasources.service.DataSourceServiceImpl; +import org.opensearch.tasks.Task; +import org.opensearch.transport.TransportService; + +@ExtendWith(MockitoExtension.class) +public class TransportPatchDataSourceActionTest { + + @Mock private TransportService transportService; + @Mock private TransportPatchDataSourceAction action; + @Mock private DataSourceServiceImpl dataSourceService; + @Mock private Task task; + @Mock private ActionListener actionListener; + + @Captor + private ArgumentCaptor patchDataSourceActionResponseArgumentCaptor; + + @Captor private ArgumentCaptor exceptionArgumentCaptor; + + @BeforeEach + public void setUp() { + action = + new TransportPatchDataSourceAction( + transportService, new ActionFilters(new HashSet<>()), dataSourceService); + } + + @Test + public void testDoExecute() { + Map dataSourceData = new HashMap<>(); + dataSourceData.put(NAME_FIELD, "test_datasource"); + dataSourceData.put(DESCRIPTION_FIELD, "test"); + + PatchDataSourceActionRequest request = new PatchDataSourceActionRequest(dataSourceData); + + action.doExecute(task, request, actionListener); + verify(dataSourceService, times(1)).patchDataSource(dataSourceData); + Mockito.verify(actionListener) + .onResponse(patchDataSourceActionResponseArgumentCaptor.capture()); + PatchDataSourceActionResponse patchDataSourceActionResponse = + patchDataSourceActionResponseArgumentCaptor.getValue(); + String responseAsJson = "\"Updated DataSource with name test_datasource\""; + Assertions.assertEquals(responseAsJson, patchDataSourceActionResponse.getResult()); + } + + @Test + public void testDoExecuteWithException() { + Map dataSourceData = new HashMap<>(); + dataSourceData.put(NAME_FIELD, "test_datasource"); + dataSourceData.put(DESCRIPTION_FIELD, "test"); + doThrow(new RuntimeException("Error")).when(dataSourceService).patchDataSource(dataSourceData); + PatchDataSourceActionRequest request = new PatchDataSourceActionRequest(dataSourceData); + action.doExecute(task, request, actionListener); + verify(dataSourceService, times(1)).patchDataSource(dataSourceData); + Mockito.verify(actionListener).onFailure(exceptionArgumentCaptor.capture()); + Exception exception = exceptionArgumentCaptor.getValue(); + Assertions.assertTrue(exception instanceof RuntimeException); + Assertions.assertEquals("Error", exception.getMessage()); + } +} diff --git a/datasources/src/test/java/org/opensearch/sql/datasources/utils/XContentParserUtilsTest.java b/datasources/src/test/java/org/opensearch/sql/datasources/utils/XContentParserUtilsTest.java index d134293456..e1e442d12b 100644 --- a/datasources/src/test/java/org/opensearch/sql/datasources/utils/XContentParserUtilsTest.java +++ b/datasources/src/test/java/org/opensearch/sql/datasources/utils/XContentParserUtilsTest.java @@ -1,6 +1,7 @@ package org.opensearch.sql.datasources.utils; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.opensearch.sql.datasources.utils.XContentParserUtils.*; import com.google.gson.Gson; import java.util.HashMap; @@ -52,6 +53,46 @@ public void testToDataSourceMetadataFromJson() { Assertions.assertEquals("prometheus_access", retrievedMetadata.getAllowedRoles().get(0)); } + @SneakyThrows + @Test + public void testToMapFromJson() { + Map dataSourceData = + Map.of( + NAME_FIELD, + "test_DS", + DESCRIPTION_FIELD, + "test", + ALLOWED_ROLES_FIELD, + List.of("all_access"), + PROPERTIES_FIELD, + Map.of("prometheus.uri", "localhost:9090"), + CONNECTOR_FIELD, + "PROMETHEUS", + RESULT_INDEX_FIELD, + ""); + + Map dataSourceDataConnectorRemoved = + Map.of( + NAME_FIELD, + "test_DS", + DESCRIPTION_FIELD, + "test", + ALLOWED_ROLES_FIELD, + List.of("all_access"), + PROPERTIES_FIELD, + Map.of("prometheus.uri", "localhost:9090"), + RESULT_INDEX_FIELD, + ""); + + Gson gson = new Gson(); + String json = gson.toJson(dataSourceData); + + Map parsedData = XContentParserUtils.toMap(json); + + Assertions.assertEquals(parsedData, dataSourceDataConnectorRemoved); + Assertions.assertEquals("test", parsedData.get(DESCRIPTION_FIELD)); + } + @SneakyThrows @Test public void testToDataSourceMetadataFromJsonWithoutName() { @@ -71,6 +112,22 @@ public void testToDataSourceMetadataFromJsonWithoutName() { Assertions.assertEquals("name and connector are required fields.", exception.getMessage()); } + @SneakyThrows + @Test + public void testToMapFromJsonWithoutName() { + Map dataSourceData = new HashMap<>(Map.of(DESCRIPTION_FIELD, "test")); + Gson gson = new Gson(); + String json = gson.toJson(dataSourceData); + + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> { + XContentParserUtils.toMap(json); + }); + Assertions.assertEquals("Name is a required field.", exception.getMessage()); + } + @SneakyThrows @Test public void testToDataSourceMetadataFromJsonWithoutConnector() { @@ -106,4 +163,21 @@ public void testToDataSourceMetadataFromJsonUsingUnknownObject() { }); Assertions.assertEquals("Unknown field: test", exception.getMessage()); } + + @SneakyThrows + @Test + public void testToMapFromJsonUsingUnknownObject() { + HashMap hashMap = new HashMap<>(); + hashMap.put("test", "test"); + Gson gson = new Gson(); + String json = gson.toJson(hashMap); + + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> { + XContentParserUtils.toMap(json); + }); + Assertions.assertEquals("Unknown field: test", exception.getMessage()); + } } diff --git a/docs/user/ppl/admin/datasources.rst b/docs/user/ppl/admin/datasources.rst index 3682153b9d..31378f6cc4 100644 --- a/docs/user/ppl/admin/datasources.rst +++ b/docs/user/ppl/admin/datasources.rst @@ -93,6 +93,19 @@ we can remove authorization and other details in case of security disabled domai "allowedRoles" : ["prometheus_access"] } +* Datasource modification PATCH API ("_plugins/_query/_datasources") :: + + PATCH https://localhost:9200/_plugins/_query/_datasources + content-type: application/json + Authorization: Basic {{username}} {{password}} + + { + "name" : "my_prometheus", + "allowedRoles" : ["all_access"] + } + + **Name is required and must exist. Connector cannot be modified and will be ignored.** + * Datasource Read GET API("_plugins/_query/_datasources/{{dataSourceName}}" :: GET https://localhost:9200/_plugins/_query/_datasources/my_prometheus @@ -114,6 +127,7 @@ Each of the datasource configuration management apis are controlled by following * cluster:admin/opensearch/datasources/create [Create POST API] * cluster:admin/opensearch/datasources/read [Get GET API] * cluster:admin/opensearch/datasources/update [Update PUT API] +* cluster:admin/opensearch/datasources/patch [Update PATCH API] * cluster:admin/opensearch/datasources/delete [Delete DELETE API] Only users mapped with roles having above actions are authorized to execute datasource management apis. 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 index 8623b9fa6f..ff36d2a887 100644 --- a/integ-test/src/test/java/org/opensearch/sql/datasource/DataSourceAPIsIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/datasource/DataSourceAPIsIT.java @@ -5,6 +5,8 @@ package org.opensearch.sql.datasource; +import static org.opensearch.sql.datasources.utils.XContentParserUtils.DESCRIPTION_FIELD; +import static org.opensearch.sql.datasources.utils.XContentParserUtils.NAME_FIELD; import static org.opensearch.sql.legacy.TestUtils.getResponseBody; import com.google.common.collect.ImmutableList; @@ -15,7 +17,9 @@ import java.io.IOException; import java.lang.reflect.Type; import java.util.ArrayList; +import java.util.HashMap; import java.util.List; +import java.util.Map; import lombok.SneakyThrows; import org.junit.AfterClass; import org.junit.Assert; @@ -136,6 +140,31 @@ public void updateDataSourceAPITest() { Assert.assertEquals( "https://randomtest.com:9090", dataSourceMetadata.getProperties().get("prometheus.uri")); Assert.assertEquals("", dataSourceMetadata.getDescription()); + + // patch datasource + Map updateDS = + new HashMap<>(Map.of(NAME_FIELD, "update_prometheus", DESCRIPTION_FIELD, "test")); + Request patchRequest = getPatchDataSourceRequest(updateDS); + Response patchResponse = client().performRequest(patchRequest); + Assert.assertEquals(200, patchResponse.getStatusLine().getStatusCode()); + String patchResponseString = getResponseBody(patchResponse); + Assert.assertEquals("\"Updated DataSource with name update_prometheus\"", patchResponseString); + + // Datasource is not immediately updated. so introducing a sleep of 2s. + Thread.sleep(2000); + + // get datasource to validate the modification. + // get datasource + Request getRequestAfterPatch = getFetchDataSourceRequest("update_prometheus"); + Response getResponseAfterPatch = client().performRequest(getRequestAfterPatch); + Assert.assertEquals(200, getResponseAfterPatch.getStatusLine().getStatusCode()); + String getResponseStringAfterPatch = getResponseBody(getResponseAfterPatch); + DataSourceMetadata dataSourceMetadataAfterPatch = + new Gson().fromJson(getResponseStringAfterPatch, DataSourceMetadata.class); + Assert.assertEquals( + "https://randomtest.com:9090", + dataSourceMetadataAfterPatch.getProperties().get("prometheus.uri")); + Assert.assertEquals("test", dataSourceMetadataAfterPatch.getDescription()); } @SneakyThrows diff --git a/integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java b/integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java index 8335ada5a7..058182f123 100644 --- a/integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java +++ b/integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java @@ -49,6 +49,7 @@ import java.nio.file.Path; import java.nio.file.Paths; import java.util.Locale; +import java.util.Map; import javax.management.MBeanServerInvocationHandler; import javax.management.ObjectName; import javax.management.remote.JMXConnector; @@ -488,6 +489,15 @@ protected static Request getUpdateDataSourceRequest(DataSourceMetadata dataSourc return request; } + protected static Request getPatchDataSourceRequest(Map dataSourceData) { + Request request = new Request("PATCH", "/_plugins/_query/_datasources"); + request.setJsonEntity(new Gson().toJson(dataSourceData)); + RequestOptions.Builder restOptionsBuilder = RequestOptions.DEFAULT.toBuilder(); + restOptionsBuilder.addHeader("Content-Type", "application/json"); + request.setOptions(restOptionsBuilder); + return request; + } + protected static Request getFetchDataSourceRequest(String name) { Request request = new Request("GET", "/_plugins/_query/_datasources" + "/" + name); if (StringUtils.isEmpty(name)) { diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java index f3fd043b63..72abc929e6 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java @@ -58,18 +58,12 @@ import org.opensearch.sql.datasources.auth.DataSourceUserAuthorizationHelperImpl; import org.opensearch.sql.datasources.encryptor.EncryptorImpl; import org.opensearch.sql.datasources.glue.GlueDataSourceFactory; -import org.opensearch.sql.datasources.model.transport.CreateDataSourceActionResponse; -import org.opensearch.sql.datasources.model.transport.DeleteDataSourceActionResponse; -import org.opensearch.sql.datasources.model.transport.GetDataSourceActionResponse; -import org.opensearch.sql.datasources.model.transport.UpdateDataSourceActionResponse; +import org.opensearch.sql.datasources.model.transport.*; import org.opensearch.sql.datasources.rest.RestDataSourceQueryAction; import org.opensearch.sql.datasources.service.DataSourceMetadataStorage; import org.opensearch.sql.datasources.service.DataSourceServiceImpl; import org.opensearch.sql.datasources.storage.OpenSearchDataSourceMetadataStorage; -import org.opensearch.sql.datasources.transport.TransportCreateDataSourceAction; -import org.opensearch.sql.datasources.transport.TransportDeleteDataSourceAction; -import org.opensearch.sql.datasources.transport.TransportGetDataSourceAction; -import org.opensearch.sql.datasources.transport.TransportUpdateDataSourceAction; +import org.opensearch.sql.datasources.transport.*; import org.opensearch.sql.legacy.esdomain.LocalClusterState; import org.opensearch.sql.legacy.executor.AsyncRestExecutor; import org.opensearch.sql.legacy.metrics.Metrics; @@ -180,6 +174,10 @@ public List getRestHandlers( new ActionType<>( TransportUpdateDataSourceAction.NAME, UpdateDataSourceActionResponse::new), TransportUpdateDataSourceAction.class), + new ActionHandler<>( + new ActionType<>( + TransportPatchDataSourceAction.NAME, PatchDataSourceActionResponse::new), + TransportPatchDataSourceAction.class), new ActionHandler<>( new ActionType<>( TransportDeleteDataSourceAction.NAME, DeleteDataSourceActionResponse::new),