forked from opensearch-project/OpenSearch
-
Notifications
You must be signed in to change notification settings - Fork 2
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Introduce Template query (opensearch-project#16818)
Introduce template query that holds the content of query which can contain placeholders and can be filled by the variables from PipelineProcessingContext produced by search processors. This allows query rewrite by the search processors. --------- Signed-off-by: Mingshi Liu <[email protected]> Co-authored-by: Michael Froh <[email protected]>
- Loading branch information
Showing
21 changed files
with
1,333 additions
and
111 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
140 changes: 140 additions & 0 deletions
140
server/src/main/java/org/opensearch/index/query/BaseQueryRewriteContext.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,140 @@ | ||
| /* | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| * | ||
| * The OpenSearch Contributors require contributions made to | ||
| * this file be licensed under the Apache-2.0 license or a | ||
| * compatible open source license. | ||
| */ | ||
|
|
||
| package org.opensearch.index.query; | ||
|
|
||
| import org.opensearch.client.Client; | ||
| import org.opensearch.common.util.concurrent.CountDown; | ||
| import org.opensearch.core.action.ActionListener; | ||
| import org.opensearch.core.common.io.stream.NamedWriteableRegistry; | ||
| import org.opensearch.core.xcontent.NamedXContentRegistry; | ||
| import org.opensearch.core.xcontent.XContentParser; | ||
|
|
||
| import java.util.ArrayList; | ||
| import java.util.List; | ||
| import java.util.function.BiConsumer; | ||
| import java.util.function.LongSupplier; | ||
|
|
||
| /** | ||
| * BaseQueryRewriteContext is a base implementation of the QueryRewriteContext interface. | ||
| * It provides core functionality for query rewriting operations in OpenSearch. | ||
| * | ||
| * This class manages the context for query rewriting, including handling of asynchronous actions, | ||
| * access to content registries, and time-related operations. | ||
| */ | ||
| public class BaseQueryRewriteContext implements QueryRewriteContext { | ||
| private final NamedXContentRegistry xContentRegistry; | ||
| private final NamedWriteableRegistry writeableRegistry; | ||
| protected final Client client; | ||
| protected final LongSupplier nowInMillis; | ||
| private final List<BiConsumer<Client, ActionListener<?>>> asyncActions = new ArrayList<>(); | ||
| private final boolean validate; | ||
|
|
||
| public BaseQueryRewriteContext( | ||
| NamedXContentRegistry xContentRegistry, | ||
| NamedWriteableRegistry writeableRegistry, | ||
| Client client, | ||
| LongSupplier nowInMillis | ||
| ) { | ||
| this(xContentRegistry, writeableRegistry, client, nowInMillis, false); | ||
| } | ||
|
|
||
| public BaseQueryRewriteContext( | ||
| NamedXContentRegistry xContentRegistry, | ||
| NamedWriteableRegistry writeableRegistry, | ||
| Client client, | ||
| LongSupplier nowInMillis, | ||
| boolean validate | ||
| ) { | ||
|
|
||
| this.xContentRegistry = xContentRegistry; | ||
| this.writeableRegistry = writeableRegistry; | ||
| this.client = client; | ||
| this.nowInMillis = nowInMillis; | ||
| this.validate = validate; | ||
| } | ||
|
|
||
| /** | ||
| * The registry used to build new {@link XContentParser}s. Contains registered named parsers needed to parse the query. | ||
| */ | ||
| public NamedXContentRegistry getXContentRegistry() { | ||
| return xContentRegistry; | ||
| } | ||
|
|
||
| /** | ||
| * Returns the time in milliseconds that is shared across all resources involved. Even across shards and nodes. | ||
| */ | ||
| public long nowInMillis() { | ||
| return nowInMillis.getAsLong(); | ||
| } | ||
|
|
||
| public NamedWriteableRegistry getWriteableRegistry() { | ||
| return writeableRegistry; | ||
| } | ||
|
|
||
| /** | ||
| * Returns an instance of {@link QueryShardContext} if available of null otherwise | ||
| */ | ||
| public QueryShardContext convertToShardContext() { | ||
| return null; | ||
| } | ||
|
|
||
| /** | ||
| * Registers an async action that must be executed before the next rewrite round in order to make progress. | ||
| * This should be used if a rewriteabel needs to fetch some external resources in order to be executed ie. a document | ||
| * from an index. | ||
| */ | ||
| public void registerAsyncAction(BiConsumer<Client, ActionListener<?>> asyncAction) { | ||
| asyncActions.add(asyncAction); | ||
| } | ||
|
|
||
| /** | ||
| * Returns <code>true</code> if there are any registered async actions. | ||
| */ | ||
| public boolean hasAsyncActions() { | ||
| return asyncActions.isEmpty() == false; | ||
| } | ||
|
|
||
| /** | ||
| * Executes all registered async actions and notifies the listener once it's done. The value that is passed to the listener is always | ||
| * <code>null</code>. The list of registered actions is cleared once this method returns. | ||
| */ | ||
| public void executeAsyncActions(ActionListener listener) { | ||
| if (asyncActions.isEmpty()) { | ||
| listener.onResponse(null); | ||
| return; | ||
| } | ||
|
|
||
| CountDown countDown = new CountDown(asyncActions.size()); | ||
| ActionListener<?> internalListener = new ActionListener() { | ||
| @Override | ||
| public void onResponse(Object o) { | ||
| if (countDown.countDown()) { | ||
| listener.onResponse(null); | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public void onFailure(Exception e) { | ||
| if (countDown.fastForward()) { | ||
| listener.onFailure(e); | ||
| } | ||
| } | ||
| }; | ||
| // make a copy to prevent concurrent modification exception | ||
| List<BiConsumer<Client, ActionListener<?>>> biConsumers = new ArrayList<>(asyncActions); | ||
| asyncActions.clear(); | ||
| for (BiConsumer<Client, ActionListener<?>> action : biConsumers) { | ||
| action.accept(client, internalListener); | ||
| } | ||
| } | ||
|
|
||
| public boolean validate() { | ||
| return validate; | ||
| } | ||
| } |
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
93 changes: 93 additions & 0 deletions
93
server/src/main/java/org/opensearch/index/query/QueryCoordinatorContext.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,93 @@ | ||
| /* | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| * | ||
| * The OpenSearch Contributors require contributions made to | ||
| * this file be licensed under the Apache-2.0 license or a | ||
| * compatible open source license. | ||
| */ | ||
|
|
||
| package org.opensearch.index.query; | ||
|
|
||
| import org.opensearch.client.Client; | ||
| import org.opensearch.common.annotation.PublicApi; | ||
| import org.opensearch.core.action.ActionListener; | ||
| import org.opensearch.core.common.io.stream.NamedWriteableRegistry; | ||
| import org.opensearch.core.xcontent.NamedXContentRegistry; | ||
| import org.opensearch.search.pipeline.PipelinedRequest; | ||
|
|
||
| import java.util.HashMap; | ||
| import java.util.Map; | ||
| import java.util.function.BiConsumer; | ||
|
|
||
| /** | ||
| * The QueryCoordinatorContext class implements the QueryRewriteContext interface and provides | ||
| * additional functionality for coordinating query rewriting in OpenSearch. | ||
| * | ||
| * This class acts as a wrapper around a QueryRewriteContext instance and a PipelinedRequest, | ||
| * allowing access to both rewrite context methods and pass over search request information. | ||
| * | ||
| * @since 2.19.0 | ||
| */ | ||
| @PublicApi(since = "2.19.0") | ||
| public class QueryCoordinatorContext implements QueryRewriteContext { | ||
| private final QueryRewriteContext rewriteContext; | ||
| private final PipelinedRequest searchRequest; | ||
|
|
||
| public QueryCoordinatorContext(QueryRewriteContext rewriteContext, PipelinedRequest searchRequest) { | ||
| this.rewriteContext = rewriteContext; | ||
| this.searchRequest = searchRequest; | ||
| } | ||
|
|
||
| @Override | ||
| public NamedXContentRegistry getXContentRegistry() { | ||
| return rewriteContext.getXContentRegistry(); | ||
| } | ||
|
|
||
| @Override | ||
| public long nowInMillis() { | ||
| return rewriteContext.nowInMillis(); | ||
| } | ||
|
|
||
| @Override | ||
| public NamedWriteableRegistry getWriteableRegistry() { | ||
| return rewriteContext.getWriteableRegistry(); | ||
| } | ||
|
|
||
| @Override | ||
| public QueryShardContext convertToShardContext() { | ||
| return rewriteContext.convertToShardContext(); | ||
| } | ||
|
|
||
| @Override | ||
| public void registerAsyncAction(BiConsumer<Client, ActionListener<?>> asyncAction) { | ||
| rewriteContext.registerAsyncAction(asyncAction); | ||
| } | ||
|
|
||
| @Override | ||
| public boolean hasAsyncActions() { | ||
| return rewriteContext.hasAsyncActions(); | ||
| } | ||
|
|
||
| @Override | ||
| public void executeAsyncActions(ActionListener listener) { | ||
| rewriteContext.executeAsyncActions(listener); | ||
| } | ||
|
|
||
| @Override | ||
| public boolean validate() { | ||
| return rewriteContext.validate(); | ||
| } | ||
|
|
||
| @Override | ||
| public QueryCoordinatorContext convertToCoordinatorContext() { | ||
| return this; | ||
| } | ||
|
|
||
| public Map<String, Object> getContextVariables() { | ||
|
|
||
| // Read from pipeline context | ||
| Map<String, Object> contextVariables = new HashMap<>(searchRequest.getPipelineProcessingContext().getAttributes()); | ||
|
|
||
| return contextVariables; | ||
| } | ||
| } |
Oops, something went wrong.