-
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.
Browse files
Browse the repository at this point in the history
* Restric cancel the scope of cancel API * Fix UT, batch query only been used for REFRESH * Update style * support cancel refresh query * fix UT * refactor code * update doc * refactor code --------- (cherry picked from commit a84c3ef) Signed-off-by: Peng Huo <[email protected]> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
- Loading branch information
1 parent
89c8234
commit 1f11fbe
Showing
12 changed files
with
328 additions
and
19 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
67 changes: 67 additions & 0 deletions
67
spark/src/main/java/org/opensearch/sql/spark/dispatcher/RefreshQueryHandler.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,67 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.sql.spark.dispatcher; | ||
|
|
||
| import org.opensearch.sql.datasource.model.DataSourceMetadata; | ||
| import org.opensearch.sql.spark.asyncquery.model.AsyncQueryJobMetadata; | ||
| import org.opensearch.sql.spark.client.EMRServerlessClient; | ||
| import org.opensearch.sql.spark.dispatcher.model.DispatchQueryContext; | ||
| import org.opensearch.sql.spark.dispatcher.model.DispatchQueryRequest; | ||
| import org.opensearch.sql.spark.dispatcher.model.DispatchQueryResponse; | ||
| import org.opensearch.sql.spark.dispatcher.model.JobType; | ||
| import org.opensearch.sql.spark.execution.statestore.StateStore; | ||
| import org.opensearch.sql.spark.flint.FlintIndexMetadata; | ||
| import org.opensearch.sql.spark.flint.FlintIndexMetadataReader; | ||
| import org.opensearch.sql.spark.flint.operation.FlintIndexOp; | ||
| import org.opensearch.sql.spark.flint.operation.FlintIndexOpCancel; | ||
| import org.opensearch.sql.spark.leasemanager.LeaseManager; | ||
| import org.opensearch.sql.spark.response.JobExecutionResponseReader; | ||
|
|
||
| /** Handle Refresh Query. */ | ||
| public class RefreshQueryHandler extends BatchQueryHandler { | ||
|
|
||
| private final FlintIndexMetadataReader flintIndexMetadataReader; | ||
| private final StateStore stateStore; | ||
| private final EMRServerlessClient emrServerlessClient; | ||
|
|
||
| public RefreshQueryHandler( | ||
| EMRServerlessClient emrServerlessClient, | ||
| JobExecutionResponseReader jobExecutionResponseReader, | ||
| FlintIndexMetadataReader flintIndexMetadataReader, | ||
| StateStore stateStore, | ||
| LeaseManager leaseManager) { | ||
| super(emrServerlessClient, jobExecutionResponseReader, leaseManager); | ||
| this.flintIndexMetadataReader = flintIndexMetadataReader; | ||
| this.stateStore = stateStore; | ||
| this.emrServerlessClient = emrServerlessClient; | ||
| } | ||
|
|
||
| @Override | ||
| public String cancelJob(AsyncQueryJobMetadata asyncQueryJobMetadata) { | ||
| String datasourceName = asyncQueryJobMetadata.getDatasourceName(); | ||
| FlintIndexMetadata indexMetadata = | ||
| flintIndexMetadataReader.getFlintIndexMetadata(asyncQueryJobMetadata.getIndexName()); | ||
| FlintIndexOp jobCancelOp = | ||
| new FlintIndexOpCancel(stateStore, datasourceName, emrServerlessClient); | ||
| jobCancelOp.apply(indexMetadata); | ||
| return asyncQueryJobMetadata.getQueryId().getId(); | ||
| } | ||
|
|
||
| @Override | ||
| public DispatchQueryResponse submit( | ||
| DispatchQueryRequest dispatchQueryRequest, DispatchQueryContext context) { | ||
| DispatchQueryResponse resp = super.submit(dispatchQueryRequest, context); | ||
| DataSourceMetadata dataSourceMetadata = context.getDataSourceMetadata(); | ||
| return new DispatchQueryResponse( | ||
| resp.getQueryId(), | ||
| resp.getJobId(), | ||
| resp.getResultIndex(), | ||
| resp.getSessionId(), | ||
| dataSourceMetadata.getName(), | ||
| JobType.BATCH, | ||
| context.getIndexQueryDetails().openSearchIndexName()); | ||
| } | ||
| } |
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
39 changes: 31 additions & 8 deletions
39
spark/src/main/java/org/opensearch/sql/spark/dispatcher/model/DispatchQueryResponse.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 |
|---|---|---|
| @@ -1,14 +1,37 @@ | ||
| package org.opensearch.sql.spark.dispatcher.model; | ||
|
|
||
| import lombok.AllArgsConstructor; | ||
| import lombok.Data; | ||
| import lombok.Getter; | ||
| import org.opensearch.sql.spark.asyncquery.model.AsyncQueryId; | ||
|
|
||
| @Data | ||
| @AllArgsConstructor | ||
| @Getter | ||
| public class DispatchQueryResponse { | ||
| private AsyncQueryId queryId; | ||
| private String jobId; | ||
| private String resultIndex; | ||
| private String sessionId; | ||
| private final AsyncQueryId queryId; | ||
| private final String jobId; | ||
| private final String resultIndex; | ||
| private final String sessionId; | ||
| private final String datasourceName; | ||
| private final JobType jobType; | ||
| private final String indexName; | ||
|
|
||
| public DispatchQueryResponse( | ||
| AsyncQueryId queryId, String jobId, String resultIndex, String sessionId) { | ||
| this(queryId, jobId, resultIndex, sessionId, null, JobType.INTERACTIVE, null); | ||
| } | ||
|
|
||
| public DispatchQueryResponse( | ||
| AsyncQueryId queryId, | ||
| String jobId, | ||
| String resultIndex, | ||
| String sessionId, | ||
| String datasourceName, | ||
| JobType jobType, | ||
| String indexName) { | ||
| this.queryId = queryId; | ||
| this.jobId = jobId; | ||
| this.resultIndex = resultIndex; | ||
| this.sessionId = sessionId; | ||
| this.datasourceName = datasourceName; | ||
| this.jobType = jobType; | ||
| this.indexName = indexName; | ||
| } | ||
| } |
Oops, something went wrong.