Disable code paths in OpenSearchDataSourceMetadataStorage

fddattal · May 28, 2024 · 023f7f5 · 023f7f5
1 parent 260ef86
commit 023f7f5
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 3 deletions.
diff --git a/...main/java/org/opensearch/sql/datasources/storage/OpenSearchDataSourceMetadataStorage.java b/...main/java/org/opensearch/sql/datasources/storage/OpenSearchDataSourceMetadataStorage.java
@@ -42,11 +42,13 @@
 import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.search.SearchHit;
 import org.opensearch.search.builder.SearchSourceBuilder;
+import org.opensearch.sql.common.setting.Settings;
 import org.opensearch.sql.datasource.model.DataSourceMetadata;
 import org.opensearch.sql.datasources.encryptor.Encryptor;
 import org.opensearch.sql.datasources.exceptions.DataSourceNotFoundException;
 import org.opensearch.sql.datasources.service.DataSourceMetadataStorage;
 import org.opensearch.sql.datasources.utils.XContentParserUtils;
+import org.opensearch.sql.opensearch.setting.OpenSearchSettings;
 
 public class OpenSearchDataSourceMetadataStorage implements DataSourceMetadataStorage {
 
@@ -61,6 +63,7 @@ public class OpenSearchDataSourceMetadataStorage implements DataSourceMetadataSt
   private final ClusterService clusterService;
 
   private final Encryptor encryptor;
+  private final OpenSearchSettings settings;
 
   /**
    * This class implements DataSourceMetadataStorage interface using OpenSearch as underlying
@@ -71,14 +74,18 @@ public class OpenSearchDataSourceMetadataStorage implements DataSourceMetadataSt
    * @param encryptor Encryptor.
    */
   public OpenSearchDataSourceMetadataStorage(
-      Client client, ClusterService clusterService, Encryptor encryptor) {
+          Client client, ClusterService clusterService, Encryptor encryptor, OpenSearchSettings settings) {
     this.client = client;
     this.clusterService = clusterService;
     this.encryptor = encryptor;
+    this.settings = settings;
   }
 
   @Override
   public List<DataSourceMetadata> getDataSourceMetadata() {
+    if (!isEnabled()) {
+      return Collections.emptyList();
+    }
     if (!this.clusterService.state().routingTable().hasIndex(DATASOURCE_INDEX_NAME)) {
       createDataSourcesIndex();
       return Collections.emptyList();
@@ -88,6 +95,9 @@ public List<DataSourceMetadata> getDataSourceMetadata() {
 
   @Override
   public Optional<DataSourceMetadata> getDataSourceMetadata(String datasourceName) {
+    if (!isEnabled()) {
+      return Optional.empty();
+    }
     if (!this.clusterService.state().routingTable().hasIndex(DATASOURCE_INDEX_NAME)) {
       createDataSourcesIndex();
       return Optional.empty();
@@ -101,6 +111,9 @@ public Optional<DataSourceMetadata> getDataSourceMetadata(String datasourceName)
 
   @Override
   public void createDataSourceMetadata(DataSourceMetadata dataSourceMetadata) {
+    if (!isEnabled()) {
+      throw new IllegalArgumentException("Data source management is disabled");
+    }
     encryptDecryptAuthenticationData(dataSourceMetadata, true);
     if (!this.clusterService.state().routingTable().hasIndex(DATASOURCE_INDEX_NAME)) {
       createDataSourcesIndex();
@@ -134,6 +147,9 @@ public void createDataSourceMetadata(DataSourceMetadata dataSourceMetadata) {
 
   @Override
   public void updateDataSourceMetadata(DataSourceMetadata dataSourceMetadata) {
+    if (!isEnabled()) {
+      throw new IllegalArgumentException("Data source management is disabled");
+    }
     encryptDecryptAuthenticationData(dataSourceMetadata, true);
     UpdateRequest updateRequest =
         new UpdateRequest(DATASOURCE_INDEX_NAME, dataSourceMetadata.getName());
@@ -163,6 +179,9 @@ public void updateDataSourceMetadata(DataSourceMetadata dataSourceMetadata) {
 
   @Override
   public void deleteDataSourceMetadata(String datasourceName) {
+    if (!isEnabled()) {
+      throw new IllegalArgumentException("Data source management is disabled");
+    }
     DeleteRequest deleteRequest = new DeleteRequest(DATASOURCE_INDEX_NAME);
     deleteRequest.id(datasourceName);
     deleteRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
@@ -302,4 +321,8 @@ private void handleSigV4PropertiesEncryptionDecryption(
         .ifPresent(list::add);
     encryptOrDecrypt(propertiesMap, isEncryption, list);
   }
+
+  private boolean isEnabled() {
+    return settings.getSettingValue(Settings.Key.DATASOURCES_ENABLED);
+  }
 }
diff --git a/.../java/org/opensearch/sql/datasources/storage/OpenSearchDataSourceMetadataStorageTest.java b/.../java/org/opensearch/sql/datasources/storage/OpenSearchDataSourceMetadataStorageTest.java
@@ -18,6 +18,7 @@
 import lombok.SneakyThrows;
 import org.apache.lucene.search.TotalHits;
 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.Answers;
@@ -41,10 +42,12 @@
 import org.opensearch.index.engine.VersionConflictEngineException;
 import org.opensearch.search.SearchHit;
 import org.opensearch.search.SearchHits;
+import org.opensearch.sql.common.setting.Settings;
 import org.opensearch.sql.datasource.model.DataSourceMetadata;
 import org.opensearch.sql.datasource.model.DataSourceType;
 import org.opensearch.sql.datasources.encryptor.Encryptor;
 import org.opensearch.sql.datasources.exceptions.DataSourceNotFoundException;
+import org.opensearch.sql.opensearch.setting.OpenSearchSettings;
 
 @ExtendWith(MockitoExtension.class)
 public class OpenSearchDataSourceMetadataStorageTest {
@@ -59,6 +62,8 @@ public class OpenSearchDataSourceMetadataStorageTest {
 
   @Mock private Encryptor encryptor;
 
+  @Mock private OpenSearchSettings openSearchSettings;
+
   @Mock(answer = Answers.RETURNS_DEEP_STUBS)
   private SearchResponse searchResponse;
 
@@ -73,6 +78,11 @@ public class OpenSearchDataSourceMetadataStorageTest {
   @Mock private SearchHit searchHit;
   @InjectMocks private OpenSearchDataSourceMetadataStorage openSearchDataSourceMetadataStorage;
 
+  @BeforeEach
+  public void configureMocks() {
+    Mockito.when(openSearchSettings.getSettingValue(ArgumentMatchers.eq(Settings.Key.DATASOURCES_ENABLED))).thenReturn(true);
+  }
+
   @SneakyThrows
   @Test
   public void testGetDataSourceMetadata() {

diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java
@@ -273,7 +273,8 @@ private DataSourceServiceImpl createDataSourceService() {
     }
     DataSourceMetadataStorage dataSourceMetadataStorage =
         new OpenSearchDataSourceMetadataStorage(
-            client, clusterService, new EncryptorImpl(masterKey));
+            client, clusterService, new EncryptorImpl(masterKey), (OpenSearchSettings) pluginSettings
+        );
     DataSourceUserAuthorizationHelper dataSourceUserAuthorizationHelper =
         new DataSourceUserAuthorizationHelperImpl(client);
     return new DataSourceServiceImpl(

diff --git a/spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceSpec.java b/spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceSpec.java
@@ -210,7 +210,7 @@ private DataSourceServiceImpl createDataSourceService() {
     String masterKey = "a57d991d9b573f75b9bba1df";
     DataSourceMetadataStorage dataSourceMetadataStorage =
         new OpenSearchDataSourceMetadataStorage(
-            client, clusterService, new EncryptorImpl(masterKey));
+            client, clusterService, new EncryptorImpl(masterKey), (OpenSearchSettings) pluginSettings);
     return new DataSourceServiceImpl(
         new ImmutableSet.Builder<DataSourceFactory>()
             .add(new GlueDataSourceFactory(pluginSettings))