forked from opensearch-project/sql
-
Notifications
You must be signed in to change notification settings - Fork 0
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
21 changed files
with
657 additions
and
5 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
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
9 changes: 9 additions & 0 deletions
9
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,9 @@ | ||
| package org.opensearch.sql.spark.client; | ||
|
|
||
| import org.opensearch.sql.spark.helper.FlintHelper; | ||
|
|
||
| public interface EmrServerlessClient { | ||
|
|
||
| String startJobRun( | ||
| String applicationId, String query, String executionRoleArn, FlintHelper flintHelper); | ||
| } |
104 changes: 104 additions & 0 deletions
104
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,104 @@ | ||
| /* | ||
| * 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.util.Set; | ||
| import org.apache.logging.log4j.LogManager; | ||
| import org.apache.logging.log4j.Logger; | ||
| 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 sparkApplicationJar; | ||
| private SparkResponse sparkResponse; | ||
| private static final Logger logger = LogManager.getLogger(EmrServerlessClientImpl.class); | ||
| private static final Set<String> terminalStates = Set.of("CANCELLED", "FAILED", "SUCCESS"); | ||
| private static final String JOB_NAME = "flint-opensearch-query"; | ||
|
|
||
| public EmrServerlessClientImpl(AWSEMRServerless emrServerless, SparkResponse sparkResponse) { | ||
| this.emrServerless = emrServerless; | ||
| this.sparkApplicationJar = SPARK_SQL_APPLICATION_JAR; | ||
| this.sparkResponse = sparkResponse; | ||
| } | ||
|
|
||
| @Override | ||
| public String startJobRun( | ||
| String applicationId, String query, String executionRoleArn, FlintHelper flint) { | ||
| StartJobRunRequest request = | ||
| new StartJobRunRequest() | ||
| .withName(JOB_NAME) | ||
| .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()); | ||
| sparkResponse.setValue(response.getJobRunId()); | ||
| return response.getJobRunId(); | ||
| } | ||
|
|
||
| public String getJobRunState(String applicationId, 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 applicationId, String jobRunId) { | ||
| // Cancel the job run | ||
| emrServerless.cancelJobRun( | ||
| new CancelJobRunRequest().withApplicationId(applicationId).withJobRunId(jobRunId)); | ||
| } | ||
| } |
21 changes: 21 additions & 0 deletions
21
spark/src/main/java/org/opensearch/sql/spark/config/SparkExecutionEngineConfig.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,21 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.sql.spark.config; | ||
|
|
||
| import com.fasterxml.jackson.annotation.JsonIgnoreProperties; | ||
| import com.google.gson.Gson; | ||
| import lombok.Data; | ||
|
|
||
| @Data | ||
| @JsonIgnoreProperties(ignoreUnknown = true) | ||
| public class SparkExecutionEngineConfig { | ||
| private String applicationId; | ||
| private String region; | ||
|
|
||
| public static SparkExecutionEngineConfig toSparkExecutionEngineConfig(String jsonString) { | ||
| return new Gson().fromJson(jsonString, SparkExecutionEngineConfig.class); | ||
| } | ||
| } |
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
52 changes: 52 additions & 0 deletions
52
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,52 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.sql.spark.dispatcher; | ||
|
|
||
| import static org.opensearch.sql.spark.data.constants.SparkConstants.FLINT_DEFAULT_AUTH; | ||
| import static org.opensearch.sql.spark.data.constants.SparkConstants.FLINT_DEFAULT_HOST; | ||
| import static org.opensearch.sql.spark.data.constants.SparkConstants.FLINT_DEFAULT_PORT; | ||
| import static org.opensearch.sql.spark.data.constants.SparkConstants.FLINT_DEFAULT_REGION; | ||
| import static org.opensearch.sql.spark.data.constants.SparkConstants.FLINT_DEFAULT_SCHEME; | ||
| import static org.opensearch.sql.spark.data.constants.SparkConstants.FLINT_INTEGRATION_JAR; | ||
|
|
||
| import lombok.AllArgsConstructor; | ||
| import org.opensearch.sql.datasource.DataSourceService; | ||
| import org.opensearch.sql.datasource.model.DataSourceMetadata; | ||
| import org.opensearch.sql.spark.client.EmrServerlessClient; | ||
| import org.opensearch.sql.spark.helper.FlintHelper; | ||
|
|
||
| @AllArgsConstructor | ||
| public class SparkQueryDispatcher { | ||
|
|
||
| private EmrServerlessClient emrServerlessClient; | ||
|
|
||
| private DataSourceService dataSourceService; | ||
|
|
||
| public String dispatch(String applicationId, String query) { | ||
| String datasourceName = getDataSourceName(query); | ||
| return emrServerlessClient.startJobRun( | ||
| applicationId, query, getExecutionRoleARN(datasourceName), getFlintHelper(datasourceName)); | ||
| } | ||
|
|
||
| private String getDataSourceName(String query) { | ||
| return "myS3glue"; | ||
| } | ||
|
|
||
| private String getExecutionRoleARN(String datasourceName) { | ||
| DataSourceMetadata dataSourceMetadata = dataSourceService.getDataSourceMetadata(datasourceName); | ||
| return dataSourceMetadata.getProperties().get("glue.auth.role_arn"); | ||
| } | ||
|
|
||
| private FlintHelper getFlintHelper(String datasourceName) { | ||
| return new FlintHelper( | ||
| FLINT_INTEGRATION_JAR, | ||
| FLINT_DEFAULT_HOST, | ||
| FLINT_DEFAULT_PORT, | ||
| FLINT_DEFAULT_SCHEME, | ||
| FLINT_DEFAULT_AUTH, | ||
| FLINT_DEFAULT_REGION); | ||
| } | ||
| } |
15 changes: 15 additions & 0 deletions
15
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,15 @@ | ||
| /* | ||
| * 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); | ||
| } |
42 changes: 42 additions & 0 deletions
42
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,42 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.sql.spark.jobs; | ||
|
|
||
| import lombok.AllArgsConstructor; | ||
| import org.opensearch.sql.common.setting.Settings; | ||
| import org.opensearch.sql.spark.config.SparkExecutionEngineConfig; | ||
| import org.opensearch.sql.spark.dispatcher.SparkQueryDispatcher; | ||
| import org.opensearch.sql.spark.jobs.model.JobMetadata; | ||
|
|
||
| @AllArgsConstructor | ||
| public class JobManagementServiceImpl implements JobManagementService { | ||
| private JobMetadataStorageService jobMetadataStorageService; | ||
| private SparkQueryDispatcher sparkQueryDispatcher; | ||
| private Settings settings; | ||
|
|
||
| @Override | ||
| public String createJob(String query) { | ||
| String sparkExecutionEngineConfigString = | ||
| settings.getSettingValue(Settings.Key.SPARK_EXECUTION_ENGINE_CONFIG); | ||
| SparkExecutionEngineConfig sparkExecutionEngineConfig = | ||
| SparkExecutionEngineConfig.toSparkExecutionEngineConfig(sparkExecutionEngineConfigString); | ||
| String jobId = | ||
| sparkQueryDispatcher.dispatch(sparkExecutionEngineConfig.getApplicationId(), query); | ||
| jobMetadataStorageService.storeJobMetadata( | ||
| new JobMetadata(jobId, sparkExecutionEngineConfig.getApplicationId())); | ||
| return jobId; | ||
| } | ||
|
|
||
| @Override | ||
| public String getJob(String jobId) { | ||
| return null; | ||
| } | ||
|
|
||
| @Override | ||
| public String cancelJob(String jobIds) { | ||
| return null; | ||
| } | ||
| } |
Oops, something went wrong.