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.
Restrict the scope of cancel API (opensearch-project#2548)
* Restric cancel the scope of cancel API Signed-off-by: Peng Huo <[email protected]> * Fix UT, batch query only been used for REFRESH Signed-off-by: Peng Huo <[email protected]> * Update style Signed-off-by: Peng Huo <[email protected]> * support cancel refresh query Signed-off-by: Peng Huo <[email protected]> * fix UT Signed-off-by: Peng Huo <[email protected]> * refactor code Signed-off-by: Peng Huo <[email protected]> * update doc Signed-off-by: Peng Huo <[email protected]> * refactor code Signed-off-by: Peng Huo <[email protected]> --------- Signed-off-by: Peng Huo <[email protected]>
- Loading branch information
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.