From 322ebf50f68a01e0f816235bb7787f85ef4ffea6 Mon Sep 17 00:00:00 2001 From: Vamsi Manohar Date: Wed, 1 Nov 2023 14:40:58 -0700 Subject: [PATCH] Block settings in sql query settings API and add more unit tests Signed-off-by: Vamsi Manohar --- .../plugin/rest/RestQuerySettingsAction.java | 17 +++++++++++++++ .../client/EmrServerlessClientImplTest.java | 21 ++++++++++++++++++- .../sql/spark/constants/TestConstants.java | 3 +++ 3 files changed, 40 insertions(+), 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 index 885c953c17..7813e0d169 100644 --- 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 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 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 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 index 319e887171..b28afe8f10 100644 --- 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,11 +8,15 @@ 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; @@ -20,13 +24,18 @@ 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 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 index 3a0d8fc56d..ca3f0ea920 100644 --- 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"; + }