-
Notifications
You must be signed in to change notification settings - Fork 141
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 <[email protected]>
- Loading branch information
Showing
13 changed files
with
299 additions
and
3 deletions.
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
7 changes: 7 additions & 0 deletions
7
spark/src/main/java/org/opensearch/sql/spark/client/EmrServerlessClient.java
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,7 @@ | ||
| package org.opensearch.sql.spark.client; | ||
|
|
||
| public interface EmrServerlessClient { | ||
|
|
||
| String startJobRun(String query); | ||
|
|
||
| } |
143 changes: 143 additions & 0 deletions
143
spark/src/main/java/org/opensearch/sql/spark/client/EmrServerlessClientImpl.java
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,143 @@ | ||
| /* | ||
| * 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_INDEX_NAME; | ||
| 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 java.util.Set; | ||
| import lombok.SneakyThrows; | ||
| 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 EmrServerlessClient { | ||
|
|
||
| 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); | ||
| private static final Set<String> terminalStates = Set.of("CANCELLED", "FAILED", "SUCCESS"); | ||
|
|
||
| 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; | ||
| } | ||
|
|
||
| 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, | ||
| SPARK_INDEX_NAME, | ||
| flint.getFlintHost(), | ||
| flint.getFlintPort(), | ||
| flint.getFlintScheme(), | ||
| flint.getFlintAuth(), | ||
| flint.getFlintRegion()) | ||
| .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); | ||
| logger.info("Job Run ID: " + response.getJobRunId()); | ||
| waitForJobToComplete(response.getJobRunId()); | ||
| return response.getJobRunId(); | ||
| } | ||
|
|
||
| @SneakyThrows | ||
| public void waitForJobToComplete(String jobRunId) { | ||
| logger.info(String.format("Waiting for job %s/%s", applicationId, jobRunId)); | ||
| int retries = 0; | ||
| while (retries < 100) { | ||
| if (!terminalStates.contains(getJobRunState(jobRunId))) { | ||
| Thread.sleep(10000); | ||
| } else { | ||
| break; | ||
| } | ||
| retries++; | ||
| } | ||
|
|
||
| if (!terminalStates.contains(getJobRunState(jobRunId))) { | ||
| throw new RuntimeException("Job was not finished after 100 retries" + jobRunId); | ||
| } | ||
| } | ||
|
|
||
| public String getJobRunState(String jobRunId) { | ||
| GetJobRunRequest request = | ||
| new GetJobRunRequest().withApplicationId(applicationId).withJobRunId(jobRunId); | ||
| GetJobRunResult response = emrServerless.getJobRun(request); | ||
| logger.info("Job Run state: " + response.getJobRun().getState()); | ||
| return response.getJobRun().getState(); | ||
| } | ||
|
|
||
| public void cancelJobRun(String jobRunId) { | ||
| // Cancel the job run | ||
| emrServerless.cancelJobRun( | ||
| new CancelJobRunRequest().withApplicationId(applicationId).withJobRunId(jobRunId)); | ||
| } | ||
|
|
||
| @Override | ||
| public String startJobRun(String query) { | ||
| String jobId = startJobRun("temp", query); | ||
| sparkResponse.setValue(jobId); | ||
| getJobRunState(jobId); | ||
| return jobId; | ||
| } | ||
| } |
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
20 changes: 20 additions & 0 deletions
20
spark/src/main/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcher.java
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,20 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.sql.spark.dispatcher; | ||
|
|
||
| import org.opensearch.sql.datasource.DataSourceService; | ||
| import org.opensearch.sql.spark.client.EmrServerlessClient; | ||
|
|
||
| public class SparkQueryDispatcher { | ||
|
|
||
| private DataSourceService dataSourceService; | ||
| private EmrServerlessClient emrServerlessClient; | ||
|
|
||
| public void dispatch(String query) { | ||
|
|
||
| } | ||
|
|
||
| } |
16 changes: 16 additions & 0 deletions
16
spark/src/main/java/org/opensearch/sql/spark/jobs/JobManagementService.java
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,16 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.sql.spark.jobs; | ||
|
|
||
| public interface JobManagementService { | ||
|
|
||
| String createJob(String query); | ||
|
|
||
| String getJob(String jobId); | ||
|
|
||
| String cancelJob(String jobIds); | ||
|
|
||
| } |
25 changes: 25 additions & 0 deletions
25
spark/src/main/java/org/opensearch/sql/spark/jobs/JobManagementServiceImpl.java
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,25 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.sql.spark.jobs; | ||
|
|
||
| public class JobManagementServiceImpl implements JobManagementService{ | ||
|
|
||
| private JobMetadataStorageService jobMetadataStorageService; | ||
| @Override | ||
| public String createJob(String query) { | ||
| return null; | ||
| } | ||
|
|
||
| @Override | ||
| public String getJob(String jobId) { | ||
| return null; | ||
| } | ||
|
|
||
| @Override | ||
| public String cancelJob(String jobIds) { | ||
| return null; | ||
| } | ||
| } |
7 changes: 7 additions & 0 deletions
7
spark/src/main/java/org/opensearch/sql/spark/jobs/JobMetadataStorageService.java
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,7 @@ | ||
| package org.opensearch.sql.spark.jobs; | ||
|
|
||
| public interface JobMetadataStorageService { | ||
|
|
||
| void createJobMetadata(); | ||
|
|
||
| } |
69 changes: 69 additions & 0 deletions
69
spark/src/test/java/org/opensearch/sql/spark/client/EmrServerlessClientImplTest.java
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,69 @@ | ||
| /* 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