Revise PR to use existing functions

Signed-off-by: Derek Ho <[email protected]>
opensearch-project · Oct 11, 2023 · 84522f1 · 84522f1
1 parent 9f4506a
commit 84522f1
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 40 deletions.
diff --git a/...urces/src/main/java/org/opensearch/sql/datasources/service/DataSourceMetadataStorage.java b/...urces/src/main/java/org/opensearch/sql/datasources/service/DataSourceMetadataStorage.java
@@ -8,7 +8,6 @@
 package org.opensearch.sql.datasources.service;
 
 import java.util.List;
-import java.util.Map;
 import java.util.Optional;
 import org.opensearch.sql.datasource.model.DataSource;
 import org.opensearch.sql.datasource.model.DataSourceMetadata;
@@ -48,13 +47,6 @@ public interface DataSourceMetadataStorage {
    */
   void updateDataSourceMetadata(DataSourceMetadata dataSourceMetadata);
 
-  /**
-   * Patches {@link DataSourceMetadata} in underlying storage.
-   *
-   * @param dataSourceData
-   */
-  void patchDataSourceMetadata(Map<String, Object> dataSourceData);
-
   /**
    * Deletes {@link DataSourceMetadata} corresponding to the datasourceName from underlying storage.
    *

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
@@ -6,7 +6,7 @@
 package org.opensearch.sql.datasources.service;
 
 import static org.opensearch.sql.analysis.DataSourceSchemaIdentifierNameResolver.DEFAULT_DATASOURCE_NAME;
-import static org.opensearch.sql.datasources.utils.XContentParserUtils.NAME_FIELD;
+import static org.opensearch.sql.datasources.utils.XContentParserUtils.*;
 
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
@@ -86,7 +86,6 @@ public void createDataSource(DataSourceMetadata metadata) {
   public void updateDataSource(DataSourceMetadata dataSourceMetadata) {
     validateDataSourceMetaData(dataSourceMetadata);
     if (!dataSourceMetadata.getName().equals(DEFAULT_DATASOURCE_NAME)) {
-      this.dataSourceLoaderCache.getOrLoadDataSource(dataSourceMetadata);
       this.dataSourceMetadataStorage.updateDataSourceMetadata(dataSourceMetadata);
     } else {
       throw new UnsupportedOperationException(
@@ -97,6 +96,10 @@ public void updateDataSource(DataSourceMetadata dataSourceMetadata) {
   @Override
   public void patchDataSource(Map<String, Object> dataSourceData) {
     if (!dataSourceData.get(NAME_FIELD).equals(DEFAULT_DATASOURCE_NAME)) {
+      DataSourceMetadata dataSourceMetadata =
+          getRawDataSourceMetadata((String) dataSourceData.get(NAME_FIELD));
+      replaceOldDatasourceMetadata(dataSourceData, dataSourceMetadata);
+      updateDataSource(dataSourceMetadata);
       this.dataSourceMetadataStorage.patchDataSourceMetadata(dataSourceData);
     } else {
       throw new UnsupportedOperationException(
@@ -140,6 +143,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<String, Object> 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<String>) dataSourceData.get(ALLOWED_ROLES_FIELD));
+          break;
+        case PROPERTIES_FIELD:
+          Map<String, String> properties = new HashMap<>(metadata.getProperties());
+          properties.putAll(((Map<String, String>) 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/...main/java/org/opensearch/sql/datasources/storage/OpenSearchDataSourceMetadataStorage.java b/...main/java/org/opensearch/sql/datasources/storage/OpenSearchDataSourceMetadataStorage.java
@@ -7,7 +7,6 @@
 
 package org.opensearch.sql.datasources.storage;
 
-import static org.opensearch.sql.datasources.utils.XContentParserUtils.NAME_FIELD;
 import static org.opensearch.sql.datasources.utils.XContentParserUtils.PROPERTIES_FIELD;
 
 import java.io.IOException;
@@ -164,35 +163,6 @@ public void updateDataSourceMetadata(DataSourceMetadata dataSourceMetadata) {
     }
   }
 
-  @Override
-  public void patchDataSourceMetadata(Map<String, Object> dataSourceData) {
-    encryptDecryptAuthenticationData(dataSourceData, true);
-    UpdateRequest updateRequest =
-        new UpdateRequest(DATASOURCE_INDEX_NAME, (String) dataSourceData.get(NAME_FIELD));
-    UpdateResponse updateResponse;
-    try (ThreadContext.StoredContext storedContext =
-        client.threadPool().getThreadContext().stashContext()) {
-      updateRequest.doc(XContentParserUtils.convertMapToXContent(dataSourceData));
-      updateRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
-      ActionFuture<UpdateResponse> updateResponseActionFuture = client.update(updateRequest);
-      updateResponse = updateResponseActionFuture.actionGet();
-    } catch (DocumentMissingException exception) {
-      throw new DataSourceNotFoundException(
-          "Datasource with name: " + dataSourceData.get(NAME_FIELD) + " doesn't exist");
-    } catch (Exception e) {
-      throw new RuntimeException(e);
-    }
-
-    if (updateResponse.getResult().equals(DocWriteResponse.Result.UPDATED)
-        || updateResponse.getResult().equals(DocWriteResponse.Result.NOOP)) {
-      LOG.debug("DatasourceMetadata : {}  successfully updated", dataSourceData.get(NAME_FIELD));
-    } else {
-      throw new RuntimeException(
-          "Saving dataSource metadata information failed with result : "
-              + updateResponse.getResult().getLowercase());
-    }
-  }
-
   @Override
   public void deleteDataSourceMetadata(String datasourceName) {
     DeleteRequest deleteRequest = new DeleteRequest(DATASOURCE_INDEX_NAME);