Block settings in sql query settings API and add more unit tests

Signed-off-by: Vamsi Manohar <reddyvam@amazon.com>
opensearch-project · Nov 1, 2023 · 322ebf5 · 322ebf5
1 parent 6cd2db0
commit 322ebf5
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 1 deletion.
diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestQuerySettingsAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestQuerySettingsAction.java
@@ -39,6 +39,7 @@ public class RestQuerySettingsAction extends BaseRestHandler {
   private static final String LEGACY_SQL_SETTINGS_PREFIX = "opendistro.sql.";
   private static final String LEGACY_PPL_SETTINGS_PREFIX = "opendistro.ppl.";
   private static final String LEGACY_COMMON_SETTINGS_PREFIX = "opendistro.query.";
+  private static final String EXECUTION_ENGINE_SETTINGS_PREFIX = "plugins.query.executionengine";
   private static final List<String> SETTINGS_PREFIX =
       ImmutableList.of(
           SQL_SETTINGS_PREFIX,
@@ -48,6 +49,11 @@ public class RestQuerySettingsAction extends BaseRestHandler {
           LEGACY_PPL_SETTINGS_PREFIX,
           LEGACY_COMMON_SETTINGS_PREFIX);
 
+  private static final List<String> DENY_LIST_PREFIX =
+      ImmutableList.of(
+          EXECUTION_ENGINE_SETTINGS_PREFIX
+      );
+
   public static final String SETTINGS_API_ENDPOINT = "/_plugins/_query/settings";
   public static final String LEGACY_SQL_SETTINGS_API_ENDPOINT = "/_opendistro/_sql/settings";
 
@@ -133,6 +139,17 @@ private Settings getAndFilterSettings(Map<String, ?> source) {
                 }
                 return true;
               });
+      //Applying DenyList Filter.
+      settingsBuilder
+          .keys()
+          .removeIf(key -> {
+            for (String prefix: DENY_LIST_PREFIX) {
+              if (key.startsWith(prefix)) {
+                return true;
+              }
+            }
+            return false;
+          });
       return settingsBuilder.build();
     } catch (IOException e) {
       throw new OpenSearchGenerationException("Failed to generate [" + source + "]", e);

diff --git a/spark/src/test/java/org/opensearch/sql/spark/client/EmrServerlessClientImplTest.java b/spark/src/test/java/org/opensearch/sql/spark/client/EmrServerlessClientImplTest.java
@@ -8,25 +8,34 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static org.opensearch.sql.spark.constants.TestConstants.DEFAULT_RESULT_INDEX;
 import static org.opensearch.sql.spark.constants.TestConstants.EMRS_APPLICATION_ID;
 import static org.opensearch.sql.spark.constants.TestConstants.EMRS_EXECUTION_ROLE;
 import static org.opensearch.sql.spark.constants.TestConstants.EMRS_JOB_NAME;
 import static org.opensearch.sql.spark.constants.TestConstants.EMR_JOB_ID;
+import static org.opensearch.sql.spark.constants.TestConstants.ENTRY_POINT_START_JAR;
 import static org.opensearch.sql.spark.constants.TestConstants.QUERY;
 import static org.opensearch.sql.spark.constants.TestConstants.SPARK_SUBMIT_PARAMETERS;
 
 import com.amazonaws.services.emrserverless.AWSEMRServerless;
 import com.amazonaws.services.emrserverless.model.CancelJobRunResult;
 import com.amazonaws.services.emrserverless.model.GetJobRunResult;
 import com.amazonaws.services.emrserverless.model.JobRun;
+import com.amazonaws.services.emrserverless.model.StartJobRunRequest;
 import com.amazonaws.services.emrserverless.model.StartJobRunResult;
 import com.amazonaws.services.emrserverless.model.ValidationException;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 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.junit.jupiter.MockitoExtension;
 import org.opensearch.sql.common.setting.Settings;
@@ -40,6 +49,8 @@ public class EmrServerlessClientImplTest {
 
   @Mock private OpenSearchSettings settings;
 
+  @Captor private ArgumentCaptor<StartJobRunRequest> startJobRunRequestArgumentCaptor;
+
   @BeforeEach
   public void setUp() {
     doReturn(emptyList()).when(settings).getSettings();
@@ -64,7 +75,15 @@ void testStartJobRun() {
             SPARK_SUBMIT_PARAMETERS,
             new HashMap<>(),
             false,
-            null));
+            DEFAULT_RESULT_INDEX));
+    verify(emrServerless, times(1))
+        .startJobRun(startJobRunRequestArgumentCaptor.capture());
+    StartJobRunRequest startJobRunRequest = startJobRunRequestArgumentCaptor.getValue();
+    Assertions.assertEquals(EMRS_APPLICATION_ID, startJobRunRequest.getApplicationId());
+    Assertions.assertEquals(EMRS_EXECUTION_ROLE, startJobRunRequest.getExecutionRoleArn());
+    Assertions.assertEquals(ENTRY_POINT_START_JAR, startJobRunRequest.getJobDriver().getSparkSubmit().getEntryPoint());
+    Assertions.assertEquals(List.of(QUERY, DEFAULT_RESULT_INDEX),
+        startJobRunRequest.getJobDriver().getSparkSubmit().getEntryPointArguments());
   }
 
   @Test

diff --git a/spark/src/test/java/org/opensearch/sql/spark/constants/TestConstants.java b/spark/src/test/java/org/opensearch/sql/spark/constants/TestConstants.java
@@ -18,4 +18,7 @@ public class TestConstants {
   public static final String TEST_CLUSTER_NAME = "TEST_CLUSTER";
   public static final String MOCK_SESSION_ID = "s-0123456";
   public static final String MOCK_STATEMENT_ID = "st-0123456";
+  public static final String ENTRY_POINT_START_JAR = "file:///home/hadoop/.ivy2/jars/org.opensearch_opensearch-spark-sql-application_2.12-0.1.0-SNAPSHOT.jar";
+  public static final String DEFAULT_RESULT_INDEX = "query_execution_result_ds1";
+
 }