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]>

Author: mingshl

CHANGELOG.md

@@ -35,6 +35,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Added a new `time` field to replace the deprecated `getTime` field in `GetStats`. ([#17009](https://github.com/opensearch-project/OpenSearch/pull/17009))
 - Improve flat_object field parsing performance by reducing two passes to a single pass ([#16297](https://github.com/opensearch-project/OpenSearch/pull/16297))
 - 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))
 
 ### Dependencies
 - Bump `com.google.cloud:google-cloud-core-http` from 2.23.0 to 2.47.0 ([#16504](https://github.com/opensearch-project/OpenSearch/pull/16504))

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/BaseQueryRewriteContext.java

@@ -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;
+    }
+}

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 GeoShapeQueryBuilder geoDisjointQuery(String name, String indexedS
     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,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;
+    }
+}