Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Signed-off-by: Vamsi Manohar <reddyvam@amazon.com>
- Loading branch information
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,110 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.sql.spark.client; | ||
|
|
||
| import static org.opensearch.sql.spark.data.constants.SparkConstants.SPARK_SQL_APPLICATION_JAR; | ||
|
|
||
| import com.amazonaws.services.emrserverless.AWSEMRServerless; | ||
| import com.amazonaws.services.emrserverless.model.CancelJobRunRequest; | ||
| import com.amazonaws.services.emrserverless.model.GetJobRunRequest; | ||
| import com.amazonaws.services.emrserverless.model.GetJobRunResult; | ||
| import com.amazonaws.services.emrserverless.model.JobDriver; | ||
| import com.amazonaws.services.emrserverless.model.SparkSubmit; | ||
| import com.amazonaws.services.emrserverless.model.StartJobRunRequest; | ||
| import com.amazonaws.services.emrserverless.model.StartJobRunResult; | ||
| import java.io.IOException; | ||
| import org.apache.logging.log4j.LogManager; | ||
| import org.apache.logging.log4j.Logger; | ||
| import org.json.JSONObject; | ||
| import org.opensearch.sql.spark.helper.FlintHelper; | ||
| import org.opensearch.sql.spark.response.SparkResponse; | ||
|
|
||
| public class EmrServerlessClientImpl implements SparkClient { | ||
|
|
||
| private final AWSEMRServerless emrServerless; | ||
| private final String applicationId; | ||
| private final String executionRoleArn; | ||
| private final FlintHelper flint; | ||
| private final String sparkApplicationJar; | ||
| private SparkResponse sparkResponse; | ||
| private static final Logger logger = LogManager.getLogger(EmrServerlessClientImpl.class); | ||
|
|
||
| public EmrServerlessClientImpl( | ||
| AWSEMRServerless emrServerless, | ||
| String applicationId, | ||
| String executionRoleArn, | ||
| FlintHelper flint, | ||
| String sparkApplicationJar, | ||
| SparkResponse sparkResponse) { | ||
| this.emrServerless = emrServerless; | ||
| this.applicationId = applicationId; | ||
| this.executionRoleArn = executionRoleArn; | ||
| this.flint = flint; | ||
| this.sparkApplicationJar = | ||
| sparkApplicationJar == null ? SPARK_SQL_APPLICATION_JAR : sparkApplicationJar; | ||
| this.sparkResponse = sparkResponse; | ||
| } | ||
|
|
||
| @Override | ||
| public JSONObject sql(String query) throws IOException { | ||
| // TODO: update/ remove for async approach | ||
| sparkResponse.setValue(startJobRun("temp", query)); | ||
| return sparkResponse.getResultFromOpensearchIndex(); | ||
|
Check warning on line 55 in spark/src/main/java/org/opensearch/sql/spark/client/EmrServerlessClientImpl.java
|
||
| } | ||
|
|
||
| public String startJobRun(String jobName, String query) { | ||
| StartJobRunRequest request = | ||
| new StartJobRunRequest() | ||
| .withName(jobName) | ||
| .withApplicationId(applicationId) | ||
| .withExecutionRoleArn(executionRoleArn) | ||
| .withJobDriver( | ||
| new JobDriver() | ||
| .withSparkSubmit( | ||
| new SparkSubmit() | ||
| .withEntryPoint(sparkApplicationJar) | ||
| .withEntryPointArguments(query) | ||
| .withSparkSubmitParameters( | ||
| "--class org.opensearch.sql.SQLJob" | ||
| + " --conf spark.driver.cores=1" | ||
| + " --conf spark.driver.memory=1g" | ||
| + " --conf spark.executor.cores=2" | ||
| + " --conf spark.executor.memory=4g" | ||
| + " --conf spark.jars=" | ||
| + flint.getFlintIntegrationJar() | ||
| + " --conf spark.datasource.flint.host=" | ||
| + flint.getFlintHost() | ||
| + " --conf spark.datasource.flint.port=" | ||
| + flint.getFlintPort() | ||
| + " --conf spark.datasource.flint.scheme=" | ||
| + flint.getFlintScheme() | ||
| + " --conf spark.datasource.flint.auth=" | ||
| + flint.getFlintAuth() | ||
| + " --conf spark.datasource.flint.region=" | ||
| + flint.getFlintRegion() | ||
| + " --conf" | ||
| + " spark.emr-serverless.driverEnv.JAVA_HOME=/usr/lib/jvm/java-17-amazon-corretto.x86_64/" | ||
| + " --conf" | ||
| + " spark.executorEnv.JAVA_HOME=/usr/lib/jvm/java-17-amazon-corretto.x86_64/" | ||
| + " --conf" | ||
| + " spark.hadoop.hive.metastore.client.factory.class=com.amazonaws.glue.catalog.metastore.AWSGlueDataCatalogHiveClientFactory"))); | ||
| StartJobRunResult response = emrServerless.startJobRun(request); | ||
| return response.getJobRunId(); | ||
| } | ||
|
|
||
| public String getJobRunState(String jobRunId) { | ||
| GetJobRunRequest request = | ||
| new GetJobRunRequest().withApplicationId(applicationId).withJobRunId(jobRunId); | ||
| GetJobRunResult response = emrServerless.getJobRun(request); | ||
| return response.getJobRun().getState(); | ||
| } | ||
|
|
||
| public void cancelJobRun(String jobRunId) { | ||
| // Cancel the job run | ||
| emrServerless.cancelJobRun( | ||
| new CancelJobRunRequest().withApplicationId(applicationId).withJobRunId(jobRunId)); | ||
| } | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,70 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.sql.spark.client; | ||
|
|
||
| import static org.mockito.ArgumentMatchers.any; | ||
| import static org.mockito.Mockito.when; | ||
| 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.QUERY; | ||
|
|
||
| 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.StartJobRunResult; | ||
| import org.junit.jupiter.api.Test; | ||
| import org.junit.jupiter.api.extension.ExtendWith; | ||
| import org.mockito.Mock; | ||
| import org.mockito.junit.jupiter.MockitoExtension; | ||
| import org.opensearch.sql.spark.helper.FlintHelper; | ||
| import org.opensearch.sql.spark.response.SparkResponse; | ||
|
|
||
| @ExtendWith(MockitoExtension.class) | ||
| public class EmrServerlessClientImplTest { | ||
| @Mock private AWSEMRServerless emrServerless; | ||
| @Mock private FlintHelper flint; | ||
| @Mock private SparkResponse sparkResponse; | ||
|
|
||
| @Test | ||
| void testStartJobRun() { | ||
| StartJobRunResult response = new StartJobRunResult(); | ||
| when(emrServerless.startJobRun(any())).thenReturn(response); | ||
|
|
||
| EmrServerlessClientImpl emrServerlessClient = | ||
| new EmrServerlessClientImpl( | ||
| emrServerless, EMRS_APPLICATION_ID, EMRS_EXECUTION_ROLE, flint, null, sparkResponse); | ||
|
|
||
| emrServerlessClient.startJobRun(EMRS_JOB_NAME, QUERY); | ||
| } | ||
|
|
||
| @Test | ||
| void testGetJobRunState() { | ||
| JobRun jobRun = new JobRun(); | ||
| jobRun.setState("Running"); | ||
| GetJobRunResult response = new GetJobRunResult(); | ||
| response.setJobRun(jobRun); | ||
| when(emrServerless.getJobRun(any())).thenReturn(response); | ||
|
|
||
| EmrServerlessClientImpl emrServerlessClient = | ||
| new EmrServerlessClientImpl( | ||
| emrServerless, EMRS_APPLICATION_ID, EMRS_EXECUTION_ROLE, flint, "", sparkResponse); | ||
|
|
||
| emrServerlessClient.getJobRunState("123"); | ||
| } | ||
|
|
||
| @Test | ||
| void testCancelJobRun() { | ||
| when(emrServerless.cancelJobRun(any())).thenReturn(new CancelJobRunResult()); | ||
|
|
||
| EmrServerlessClientImpl emrServerlessClient = | ||
| new EmrServerlessClientImpl( | ||
| emrServerless, EMRS_APPLICATION_ID, EMRS_EXECUTION_ROLE, flint, null, sparkResponse); | ||
|
|
||
| emrServerlessClient.cancelJobRun("123"); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters