Skip to content

Commit 1c7f719

Browse files
msfrohgithub-actions[bot]
msfroh
and
github-actions[bot]
authored
[Backport 2.x] Introduce Template query (opensearch-project#16818) (opensearch-project#17142)
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. Backport 2.x note: Instead of extracting an interface from QueryRewriteContext, we can keep it as a base class and subclass it for QueryCoordinatorContext. This avoids a breaking API change on the 2.x line. --------- Signed-off-by: Mingshi Liu <mingshl@amazon.com> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Signed-off-by: Michael Froh <froh@amazon.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
1 parent de6b87b commit 1c7f719

File tree

13 files changed

+1137
-26
lines changed

13 files changed

+1137
-26
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
3131
- Added a new `time` field to replace the deprecated `getTime` field in `GetStats`. ([#17009](https://github.com/opensearch-project/OpenSearch/pull/17009))
3232
- Improve flat_object field parsing performance by reducing two passes to a single pass ([#16297](https://github.com/opensearch-project/OpenSearch/pull/16297))
3333
- Improve performance of the bitmap filtering([#16936](https://github.com/opensearch-project/OpenSearch/pull/16936/))
34+
- Introduce Template query ([#16818](https://github.com/opensearch-project/OpenSearch/pull/16818))
3435
- Added ability to retrieve value from DocValues in a flat_object filed([#16802](https://github.com/opensearch-project/OpenSearch/pull/16802))
3536
- Added new Setting property UnmodifiableOnRestore to prevent updating settings on restore snapshot ([#16957](https://github.com/opensearch-project/OpenSearch/pull/16957))
3637
- Introduce Template query ([#16818](https://github.com/opensearch-project/OpenSearch/pull/16818))

‎server/build.gradle‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,6 @@ dependencies {
7070
api project(":libs:opensearch-telemetry")
7171
api project(":libs:opensearch-task-commons")
7272

73-
7473
compileOnly project(':libs:opensearch-plugin-classloader')
7574
testRuntimeOnly project(':libs:opensearch-plugin-classloader')
7675

‎server/src/main/java/org/opensearch/action/search/TransportSearchAction.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -476,7 +476,7 @@ private void executeRequest(
476476
} else {
477477
Rewriteable.rewriteAndFetch(
478478
sr.source(),
479-
searchService.getRewriteContext(timeProvider::getAbsoluteStartMillis),
479+
searchService.getRewriteContext(timeProvider::getAbsoluteStartMillis, searchRequest),
480480
rewriteListener
481481
);
482482
}

‎server/src/main/java/org/opensearch/index/query/QueryBuilders.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@
5050
import java.io.IOException;
5151
import java.util.Collection;
5252
import java.util.List;
53+
import java.util.Map;
5354

5455
/**
5556
* Utility class to create search queries.
@@ -780,4 +781,13 @@ public static GeoShapeQueryBuilder geoDisjointQuery(String name, String indexedS
780781
public static ExistsQueryBuilder existsQuery(String name) {
781782
return new ExistsQueryBuilder(name);
782783
}
784+
785+
/**
786+
* A query that contains a template with holder that should be resolved by search processors
787+
*
788+
* @param content The content of the template
789+
*/
790+
public static TemplateQueryBuilder templateQuery(Map<String, Object> content) {
791+
return new TemplateQueryBuilder(content);
792+
}
783793
}

‎server/src/main/java/org/opensearch/index/query/QueryCoordinatorContext.java‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.index.query;
10+
11+
import org.opensearch.common.annotation.PublicApi;
12+
import org.opensearch.search.pipeline.PipelinedRequest;
13+
14+
import java.util.HashMap;
15+
import java.util.Map;
16+
17+
/**
18+
* The QueryCoordinatorContext class implements the QueryRewriteContext interface and provides
19+
* additional functionality for coordinating query rewriting in OpenSearch.
20+
*
21+
* This class acts as a wrapper around a QueryRewriteContext instance and a PipelinedRequest,
22+
* allowing access to both rewrite context methods and pass over search request information.
23+
*
24+
* @since 2.19.0
25+
*/
26+
@PublicApi(since = "2.19.0")
27+
public class QueryCoordinatorContext extends QueryRewriteContext {
28+
private final PipelinedRequest searchRequest;
29+
30+
public QueryCoordinatorContext(QueryRewriteContext parent, PipelinedRequest pipelinedRequest) {
31+
super(parent.getXContentRegistry(), parent.getWriteableRegistry(), parent.client, parent.nowInMillis, parent.validate());
32+
this.searchRequest = pipelinedRequest;
33+
}
34+
35+
@Override
36+
public QueryCoordinatorContext convertToCoordinatorContext() {
37+
return this;
38+
}
39+
40+
public Map<String, Object> getContextVariables() {
41+
42+
// Read from pipeline context
43+
Map<String, Object> contextVariables = new HashMap<>(searchRequest.getPipelineProcessingContext().getAttributes());
44+
45+
return contextVariables;
46+
}
47+
}

‎server/src/main/java/org/opensearch/index/query/QueryRewriteContext.java‎

Lines changed: 29 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -101,12 +101,20 @@ public NamedWriteableRegistry getWriteableRegistry() {
101101
}
102102

103103
/**
104-
* Returns an instance of {@link QueryShardContext} if available of null otherwise
104+
* Returns an instance of {@link QueryShardContext} if available or null otherwise
105105
*/
106106
public QueryShardContext convertToShardContext() {
107107
return null;
108108
}
109109

110+
/**
111+
* Returns an instance of {@link QueryCoordinatorContext} if available or null otherwise
112+
* @return
113+
*/
114+
public QueryCoordinatorContext convertToCoordinatorContext() {
115+
return null;
116+
}
117+
110118
/**
111119
* Registers an async action that must be executed before the next rewrite round in order to make progress.
112120
* This should be used if a rewriteabel needs to fetch some external resources in order to be executed ie. a document
@@ -130,29 +138,30 @@ public boolean hasAsyncActions() {
130138
public void executeAsyncActions(ActionListener listener) {
131139
if (asyncActions.isEmpty()) {
132140
listener.onResponse(null);
133-
} else {
134-
CountDown countDown = new CountDown(asyncActions.size());
135-
ActionListener<?> internalListener = new ActionListener() {
136-
@Override
137-
public void onResponse(Object o) {
138-
if (countDown.countDown()) {
139-
listener.onResponse(null);
140-
}
141+
return;
142+
}
143+
144+
CountDown countDown = new CountDown(asyncActions.size());
145+
ActionListener<?> internalListener = new ActionListener() {
146+
@Override
147+
public void onResponse(Object o) {
148+
if (countDown.countDown()) {
149+
listener.onResponse(null);
141150
}
151+
}
142152

143-
@Override
144-
public void onFailure(Exception e) {
145-
if (countDown.fastForward()) {
146-
listener.onFailure(e);
147-
}
153+
@Override
154+
public void onFailure(Exception e) {
155+
if (countDown.fastForward()) {
156+
listener.onFailure(e);
148157
}
149-
};
150-
// make a copy to prevent concurrent modification exception
151-
List<BiConsumer<Client, ActionListener<?>>> biConsumers = new ArrayList<>(asyncActions);
152-
asyncActions.clear();
153-
for (BiConsumer<Client, ActionListener<?>> action : biConsumers) {
154-
action.accept(client, internalListener);
155158
}
159+
};
160+
// make a copy to prevent concurrent modification exception
161+
List<BiConsumer<Client, ActionListener<?>>> biConsumers = new ArrayList<>(asyncActions);
162+
asyncActions.clear();
163+
for (BiConsumer<Client, ActionListener<?>> action : biConsumers) {
164+
action.accept(client, internalListener);
156165
}
157166
}
158167

‎server/src/main/java/org/opensearch/index/query/TemplateQueryBuilder.java‎

Lines changed: 198 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,198 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.index.query;
10+
11+
import org.apache.lucene.search.Query;
12+
import org.opensearch.common.xcontent.LoggingDeprecationHandler;
13+
import org.opensearch.common.xcontent.XContentFactory;
14+
import org.opensearch.common.xcontent.XContentType;
15+
import org.opensearch.common.xcontent.json.JsonXContent;
16+
import org.opensearch.core.common.io.stream.StreamInput;
17+
import org.opensearch.core.common.io.stream.StreamOutput;
18+
import org.opensearch.core.xcontent.XContentBuilder;
19+
import org.opensearch.core.xcontent.XContentParser;
20+
21+
import java.io.IOException;
22+
import java.util.Map;
23+
import java.util.Objects;
24+
import java.util.stream.Collectors;
25+
26+
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
27+
28+
/**
29+
* A query builder that constructs a query based on a template and context variables.
30+
* This query is designed to be rewritten with variables from search processors.
31+
*/
32+
33+
public class TemplateQueryBuilder extends AbstractQueryBuilder<TemplateQueryBuilder> {
34+
public static final String NAME = "template";
35+
public static final String queryName = "template";
36+
private final Map<String, Object> content;
37+
38+
/**
39+
* Constructs a new TemplateQueryBuilder with the given content.
40+
*
41+
* @param content The template content as a map.
42+
*/
43+
public TemplateQueryBuilder(Map<String, Object> content) {
44+
this.content = content;
45+
}
46+
47+
/**
48+
* Creates a TemplateQueryBuilder from XContent.
49+
*
50+
* @param parser The XContentParser to read from.
51+
* @return A new TemplateQueryBuilder instance.
52+
* @throws IOException If there's an error parsing the content.
53+
*/
54+
public static TemplateQueryBuilder fromXContent(XContentParser parser) throws IOException {
55+
return new TemplateQueryBuilder(parser.map());
56+
}
57+
58+
/**
59+
* Constructs a TemplateQueryBuilder from a stream input.
60+
*
61+
* @param in The StreamInput to read from.
62+
* @throws IOException If there's an error reading from the stream.
63+
*/
64+
public TemplateQueryBuilder(StreamInput in) throws IOException {
65+
super(in);
66+
this.content = in.readMap();
67+
}
68+
69+
@Override
70+
protected void doWriteTo(StreamOutput out) throws IOException {
71+
out.writeMap(content);
72+
}
73+
74+
@Override
75+
protected void doXContent(XContentBuilder builder, Params params) throws IOException {
76+
builder.field(NAME, content);
77+
}
78+
79+
@Override
80+
protected Query doToQuery(QueryShardContext context) throws IOException {
81+
throw new IllegalStateException(
82+
"Template queries cannot be converted directly to a query. Template Query must be rewritten first during doRewrite."
83+
);
84+
}
85+
86+
@Override
87+
protected boolean doEquals(TemplateQueryBuilder other) {
88+
return Objects.equals(this.content, other.content);
89+
}
90+
91+
@Override
92+
protected int doHashCode() {
93+
return Objects.hash(content);
94+
}
95+
96+
@Override
97+
public String getWriteableName() {
98+
return NAME;
99+
}
100+
101+
/**
102+
* Gets the content of this template query.
103+
*
104+
* @return The template content as a map.
105+
*/
106+
public Map<String, Object> getContent() {
107+
return content;
108+
}
109+
110+
/**
111+
* Rewrites the template query by substituting variables from the context.
112+
*
113+
* @param queryRewriteContext The context for query rewriting.
114+
* @return A rewritten QueryBuilder.
115+
* @throws IOException If there's an error during rewriting.
116+
*/
117+
@Override
118+
protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws IOException {
119+
// the queryRewrite is expected at QueryCoordinator level
120+
QueryCoordinatorContext queryCoordinatorContext = queryRewriteContext.convertToCoordinatorContext();
121+
if (queryCoordinatorContext == null) {
122+
throw new IllegalStateException(
123+
"Template Query must be rewritten at the coordinator node. Rewriting at shard level is not supported."
124+
);
125+
}
126+
127+
Map<String, Object> contextVariables = queryCoordinatorContext.getContextVariables();
128+
String queryString;
129+
130+
try (XContentBuilder builder = XContentFactory.jsonBuilder()) {
131+
builder.map(this.content);
132+
queryString = builder.toString();
133+
}
134+
135+
// Convert Map<String, Object> to Map<String, String> with proper JSON escaping
136+
Map<String, String> variablesMap = null;
137+
if (contextVariables != null) {
138+
variablesMap = contextVariables.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, entry -> {
139+
try {
140+
return JsonXContent.contentBuilder().value(entry.getValue()).toString();
141+
} catch (IOException e) {
142+
throw new RuntimeException("Error converting contextVariables to JSON string", e);
143+
}
144+
}));
145+
}
146+
String newQueryContent = replaceVariables(queryString, variablesMap);
147+
148+
try {
149+
XContentParser parser = XContentType.JSON.xContent()
150+
.createParser(queryCoordinatorContext.getXContentRegistry(), LoggingDeprecationHandler.INSTANCE, newQueryContent);
151+
152+
ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
153+
154+
QueryBuilder newQueryBuilder = parseInnerQueryBuilder(parser);
155+
156+
return newQueryBuilder;
157+
158+
} catch (Exception e) {
159+
throw new IllegalArgumentException("Failed to rewrite template query: " + newQueryContent, e);
160+
}
161+
}
162+
163+
private String replaceVariables(String template, Map<String, String> variables) {
164+
if (template == null || template.equals("null")) {
165+
throw new IllegalArgumentException("Template string cannot be null. A valid template must be provided.");
166+
}
167+
if (template.isEmpty() || template.equals("{}")) {
168+
throw new IllegalArgumentException("Template string cannot be empty. A valid template must be provided.");
169+
}
170+
if (variables == null || variables.isEmpty()) {
171+
return template;
172+
}
173+
174+
StringBuilder result = new StringBuilder();
175+
int start = 0;
176+
while (true) {
177+
int startVar = template.indexOf("\"${", start);
178+
if (startVar == -1) {
179+
result.append(template.substring(start));
180+
break;
181+
}
182+
result.append(template, start, startVar);
183+
int endVar = template.indexOf("}\"", startVar);
184+
if (endVar == -1) {
185+
throw new IllegalArgumentException("Unclosed variable in template: " + template.substring(startVar));
186+
}
187+
String varName = template.substring(startVar + 3, endVar);
188+
String replacement = variables.get(varName);
189+
if (replacement == null) {
190+
throw new IllegalArgumentException("Variable not found: " + varName);
191+
}
192+
result.append(replacement);
193+
start = endVar + 2;
194+
}
195+
return result.toString();
196+
}
197+
198+
}

‎server/src/main/java/org/opensearch/search/SearchModule.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@
8686
import org.opensearch.index.query.SpanOrQueryBuilder;
8787
import org.opensearch.index.query.SpanTermQueryBuilder;
8888
import org.opensearch.index.query.SpanWithinQueryBuilder;
89+
import org.opensearch.index.query.TemplateQueryBuilder;
8990
import org.opensearch.index.query.TermQueryBuilder;
9091
import org.opensearch.index.query.TermsQueryBuilder;
9192
import org.opensearch.index.query.TermsSetQueryBuilder;
@@ -1208,7 +1209,7 @@ private void registerQueryParsers(List<SearchPlugin> plugins) {
12081209
registerQuery(
12091210
new QuerySpec<>(MatchBoolPrefixQueryBuilder.NAME, MatchBoolPrefixQueryBuilder::new, MatchBoolPrefixQueryBuilder::fromXContent)
12101211
);
1211-
1212+
registerQuery(new QuerySpec<>(TemplateQueryBuilder.NAME, TemplateQueryBuilder::new, TemplateQueryBuilder::fromXContent));
12121213
if (ShapesAvailability.JTS_AVAILABLE && ShapesAvailability.SPATIAL4J_AVAILABLE) {
12131214
registerQuery(new QuerySpec<>(GeoShapeQueryBuilder.NAME, GeoShapeQueryBuilder::new, GeoShapeQueryBuilder::fromXContent));
12141215
}

0 commit comments

Comments
 (0)