Backport/backport 16818 to 2.x take2

CHANGELOG.md

@@ -29,6 +29,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Support searching from doc_value using termQueryCaseInsensitive/termQuery in flat_object/keyword field([#16974](https://github.com/opensearch-project/OpenSearch/pull/16974/))
 - Added a new `time` field to replace the deprecated `getTime` field in `GetStats`. ([#17009](https://github.com/opensearch-project/OpenSearch/pull/17009))
 - Improve performance of the bitmap filtering([#16936](https://github.com/opensearch-project/OpenSearch/pull/16936/))
+- Introduce Template query ([#16818](https://github.com/opensearch-project/OpenSearch/pull/16818))
 - Added ability to retrieve value from DocValues in a flat_object filed([#16802](https://github.com/opensearch-project/OpenSearch/pull/16802))
 - Added new Setting property UnmodifiableOnRestore to prevent updating settings on restore snapshot ([#16957](https://github.com/opensearch-project/OpenSearch/pull/16957))
 - Introduce Template query ([#16818](https://github.com/opensearch-project/OpenSearch/pull/16818))

server/build.gradle

@@ -70,7 +70,6 @@ dependencies {
   api project(":libs:opensearch-telemetry")
   api project(":libs:opensearch-task-commons")
 
-
   compileOnly project(':libs:opensearch-plugin-classloader')
   testRuntimeOnly project(':libs:opensearch-plugin-classloader')
 

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

@@ -476,7 +476,7 @@ private void executeRequest(
                 } else {
                     Rewriteable.rewriteAndFetch(
                         sr.source(),
-                        searchService.getRewriteContext(timeProvider::getAbsoluteStartMillis),
+                        searchService.getRewriteContext(timeProvider::getAbsoluteStartMillis, searchRequest),
                         rewriteListener
                     );
                 }

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

@@ -50,6 +50,7 @@
 import java.io.IOException;
 import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 
 /**
  * Utility class to create search queries.
@@ -780,4 +781,13 @@
     public static ExistsQueryBuilder existsQuery(String name) {
         return new ExistsQueryBuilder(name);
     }
+
+    /**
+     *  A query that contains a template with holder that should be resolved by search processors
+     *
+     * @param content The content of the template
+     */
+    public static TemplateQueryBuilder templateQuery(Map<String, Object> content) {
+        return new TemplateQueryBuilder(content);
+    }
 }

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

@@ -0,0 +1,47 @@
+/*
+ * 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.common.annotation.PublicApi;
+import org.opensearch.search.pipeline.PipelinedRequest;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * 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 extends QueryRewriteContext {
+    private final PipelinedRequest searchRequest;
+
+    public QueryCoordinatorContext(QueryRewriteContext parent, PipelinedRequest pipelinedRequest) {
+        super(parent.getXContentRegistry(), parent.getWriteableRegistry(), parent.client, parent.nowInMillis, parent.validate());
+        this.searchRequest = pipelinedRequest;
+    }
+
+    @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;
+    }
+}

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

@@ -101,12 +101,20 @@
     }
 
     /**
-     * Returns an instance of {@link QueryShardContext} if available of null otherwise
+     * Returns an instance of {@link QueryShardContext} if available or null otherwise
      */
     public QueryShardContext convertToShardContext() {
         return null;
     }
 
+    /**
+     * Returns an instance of {@link QueryCoordinatorContext} if available or null otherwise
+     * @return
+     */
+    public QueryCoordinatorContext convertToCoordinatorContext() {
+        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
@@ -130,29 +138,30 @@
     public void executeAsyncActions(ActionListener listener) {
         if (asyncActions.isEmpty()) {
             listener.onResponse(null);
-        } else {
-            CountDown countDown = new CountDown(asyncActions.size());
-            ActionListener<?> internalListener = new ActionListener() {
-                @Override
-                public void onResponse(Object o) {
-                    if (countDown.countDown()) {
-                        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);
-                    }
+            @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);
             }
+        };
+        // 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);
         }
     }
 

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

@@ -0,0 +1,198 @@
+/*
+ * 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.apache.lucene.search.Query;
+import org.opensearch.common.xcontent.LoggingDeprecationHandler;
+import org.opensearch.common.xcontent.XContentFactory;
+import org.opensearch.common.xcontent.XContentType;
+import org.opensearch.common.xcontent.json.JsonXContent;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.core.xcontent.XContentParser;
+
+import java.io.IOException;
+import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Collectors;
+
+import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
+
+/**
+ * A query builder that constructs a query based on a template and context variables.
+ * This query is designed to be rewritten with variables from search processors.
+ */
+
+public class TemplateQueryBuilder extends AbstractQueryBuilder<TemplateQueryBuilder> {
+    public static final String NAME = "template";
+    public static final String queryName = "template";
+    private final Map<String, Object> content;
+
+    /**
+     * Constructs a new TemplateQueryBuilder with the given content.
+     *
+     * @param content The template content as a map.
+     */
+    public TemplateQueryBuilder(Map<String, Object> content) {
+        this.content = content;
+    }
+
+    /**
+     * Creates a TemplateQueryBuilder from XContent.
+     *
+     * @param parser The XContentParser to read from.
+     * @return A new TemplateQueryBuilder instance.
+     * @throws IOException If there's an error parsing the content.
+     */
+    public static TemplateQueryBuilder fromXContent(XContentParser parser) throws IOException {
+        return new TemplateQueryBuilder(parser.map());
+    }
+
+    /**
+     * Constructs a TemplateQueryBuilder from a stream input.
+     *
+     * @param in The StreamInput to read from.
+     * @throws IOException If there's an error reading from the stream.
+     */
+    public TemplateQueryBuilder(StreamInput in) throws IOException {
+        super(in);
+        this.content = in.readMap();
+    }
+
+    @Override
+    protected void doWriteTo(StreamOutput out) throws IOException {
+        out.writeMap(content);
+    }
+
+    @Override
+    protected void doXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.field(NAME, content);
+    }
+
+    @Override
+    protected Query doToQuery(QueryShardContext context) throws IOException {
+        throw new IllegalStateException(
+            "Template queries cannot be converted directly to a query. Template Query must be rewritten first during doRewrite."
+        );
+    }
+
+    @Override
+    protected boolean doEquals(TemplateQueryBuilder other) {
+        return Objects.equals(this.content, other.content);
+    }
+
+    @Override
+    protected int doHashCode() {
+        return Objects.hash(content);
+    }
+
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
+
+    /**
+     * Gets the content of this template query.
+     *
+     * @return The template content as a map.
+     */
+    public Map<String, Object> getContent() {
+        return content;
+    }
+
+    /**
+     * Rewrites the template query by substituting variables from the context.
+     *
+     * @param queryRewriteContext The context for query rewriting.
+     * @return A rewritten QueryBuilder.
+     * @throws IOException If there's an error during rewriting.
+     */
+    @Override
+    protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws IOException {
+        // the queryRewrite is expected at QueryCoordinator level
+        QueryCoordinatorContext queryCoordinatorContext = queryRewriteContext.convertToCoordinatorContext();
+        if (queryCoordinatorContext == null) {
+            throw new IllegalStateException(
+                "Template Query must be rewritten at the coordinator node. Rewriting at shard level is not supported."
+            );
+        }
+
+        Map<String, Object> contextVariables = queryCoordinatorContext.getContextVariables();
+        String queryString;
+
+        try (XContentBuilder builder = XContentFactory.jsonBuilder()) {
+            builder.map(this.content);
+            queryString = builder.toString();
+        }
+
+        // Convert Map<String, Object> to Map<String, String> with proper JSON escaping
+        Map<String, String> variablesMap = null;
+        if (contextVariables != null) {
+            variablesMap = contextVariables.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, entry -> {
+                try {
+                    return JsonXContent.contentBuilder().value(entry.getValue()).toString();
+                } catch (IOException e) {
+                    throw new RuntimeException("Error converting contextVariables to JSON string", e);
+                }
+            }));
+        }
+        String newQueryContent = replaceVariables(queryString, variablesMap);
+
+        try {
+            XContentParser parser = XContentType.JSON.xContent()
+                .createParser(queryCoordinatorContext.getXContentRegistry(), LoggingDeprecationHandler.INSTANCE, newQueryContent);
+
+            ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
+
+            QueryBuilder newQueryBuilder = parseInnerQueryBuilder(parser);
+
+            return newQueryBuilder;
+
+        } catch (Exception e) {
+            throw new IllegalArgumentException("Failed to rewrite template query: " + newQueryContent, e);
+        }
+    }
+
+    private String replaceVariables(String template, Map<String, String> variables) {
+        if (template == null || template.equals("null")) {
+            throw new IllegalArgumentException("Template string cannot be null. A valid template must be provided.");
+        }
+        if (template.isEmpty() || template.equals("{}")) {
+            throw new IllegalArgumentException("Template string cannot be empty. A valid template must be provided.");
+        }
+        if (variables == null || variables.isEmpty()) {
+            return template;
+        }
+
+        StringBuilder result = new StringBuilder();
+        int start = 0;
+        while (true) {
+            int startVar = template.indexOf("\"${", start);
+            if (startVar == -1) {
+                result.append(template.substring(start));
+                break;
+            }
+            result.append(template, start, startVar);
+            int endVar = template.indexOf("}\"", startVar);
+            if (endVar == -1) {
+                throw new IllegalArgumentException("Unclosed variable in template: " + template.substring(startVar));
+            }
+            String varName = template.substring(startVar + 3, endVar);
+            String replacement = variables.get(varName);
+            if (replacement == null) {
+                throw new IllegalArgumentException("Variable not found: " + varName);
+            }
+            result.append(replacement);
+            start = endVar + 2;
+        }
+        return result.toString();
+    }
+
+}

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

@@ -86,6 +86,7 @@
 import org.opensearch.index.query.SpanOrQueryBuilder;
 import org.opensearch.index.query.SpanTermQueryBuilder;
 import org.opensearch.index.query.SpanWithinQueryBuilder;
+import org.opensearch.index.query.TemplateQueryBuilder;
 import org.opensearch.index.query.TermQueryBuilder;
 import org.opensearch.index.query.TermsQueryBuilder;
 import org.opensearch.index.query.TermsSetQueryBuilder;
@@ -1208,7 +1209,7 @@ private void registerQueryParsers(List<SearchPlugin> plugins) {
         registerQuery(
             new QuerySpec<>(MatchBoolPrefixQueryBuilder.NAME, MatchBoolPrefixQueryBuilder::new, MatchBoolPrefixQueryBuilder::fromXContent)
         );
-
+        registerQuery(new QuerySpec<>(TemplateQueryBuilder.NAME, TemplateQueryBuilder::new, TemplateQueryBuilder::fromXContent));
         if (ShapesAvailability.JTS_AVAILABLE && ShapesAvailability.SPATIAL4J_AVAILABLE) {
             registerQuery(new QuerySpec<>(GeoShapeQueryBuilder.NAME, GeoShapeQueryBuilder::new, GeoShapeQueryBuilder::fromXContent));
         }