Add Template Query

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

plugins/ingest-attachment/build.gradle

@@ -99,7 +99,7 @@ dependencies {
   api "org.apache.james:apache-mime4j-core:${versions.mime4j}"
   api "org.apache.james:apache-mime4j-dom:${versions.mime4j}"
   // EPUB books
-  api "org.apache.commons:commons-lang3:${versions.commonslang}"
+  api "org.apache.commons:commons-lang3:3.17.0"
   // Microsoft Word files with visio diagrams
   api 'org.apache.commons:commons-math3:3.6.1'
   // POIs dependency

plugins/repository-azure/build.gradle

@@ -83,7 +83,7 @@ dependencies {
   api 'org.codehaus.woodstox:stax2-api:4.2.2'
   implementation "com.fasterxml.woodstox:woodstox-core:${versions.woodstox}"
   runtimeOnly "com.google.guava:guava:${versions.guava}"
-  api "org.apache.commons:commons-lang3:${versions.commonslang}"
+  api "org.apache.commons:commons-lang3:3.17.0"
   testImplementation project(':test:fixtures:azure-fixture')
 }
 

server/build.gradle

@@ -70,7 +70,9 @@ dependencies {
   api project(":libs:opensearch-telemetry")
   api project(":libs:opensearch-task-commons")
 
-
+  // TODO: need to move to proper packages under libs
+  implementation group: 'org.apache.commons', name: 'commons-text', version: '1.10.0'
+  api 'org.apache.commons:commons-lang3:3.17.0'
   compileOnly project(':libs:opensearch-plugin-classloader')
   testRuntimeOnly project(':libs:opensearch-plugin-classloader')
 

server/licenses/commons-text-LICENSE.txt

@@ -0,0 +1,8 @@
+This copy of commons-text is licensed under the
+Apache (Software) License, version 2.0 ("the License").
+See the License for details about distribution rights, and the
+specific rights regarding derivate works.
+
+You may obtain a copy of the License at:
+
+http://www.apache.org/licenses/LICENSE-2.0

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

@@ -20,6 +20,9 @@
 import java.util.function.BiConsumer;
 import java.util.function.LongSupplier;
 
+/**
+ * BaseQueryRewriteContext
+ */
 public class BaseQueryRewriteContext implements QueryRewriteContext {
     private final NamedXContentRegistry xContentRegistry;
     private final NamedWriteableRegistry writeableRegistry;

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

@@ -9,18 +9,24 @@
 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;
 
+/**
+ * QueryCoordinatorContext
+ */
+@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;
@@ -71,13 +77,13 @@ public QueryCoordinatorContext convertToCoordinatorContext() {
         return this;
     }
 
-    public Object getContextVariable(String variableName) {
-        // Read from request search exts
+    public Map<String, Object> getContextVariables() {
+
+        Map<String, Object> contextVariables = new HashMap<>();
 
         // Read from pipeline context
-        Object val = searchRequest.getPipelineProcessingContext().getAttribute(variableName);
-        if (val != null) {
-            return val;
-        }
+        contextVariables.putAll(searchRequest.getPipelineProcessingContext().getAttributes());
+
+        return contextVariables;
     }
 }

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

@@ -33,16 +33,12 @@
 
 import org.opensearch.client.Client;
 import org.opensearch.common.annotation.PublicApi;
-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;
 
 /**
  * Context object used to rewrite {@link QueryBuilder} instances into simplified version.

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

@@ -0,0 +1,158 @@
+/*
+ * 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.commons.text.StringSubstitutor;
+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 query should run with a ml_inference request processor");
+    }
+
+    @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 queryShardContext The context for query rewriting.
+     * @return A rewritten QueryBuilder.
+     * @throws IOException If there's an error during rewriting.
+     */
+    @Override
+    protected QueryBuilder doRewrite(QueryRewriteContext queryShardContext) throws IOException {
+        if (!(queryShardContext instanceof QueryCoordinatorContext)) {
+            throw new IllegalStateException("Template query needs to be resolved with variables from search processors.");
+        }
+
+        QueryCoordinatorContext queryCoordinateContext = (QueryCoordinatorContext) queryShardContext;
+
+        Map<String, Object> contextVariables = queryCoordinateContext.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 = 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);
+            }
+        }));
+
+        StringSubstitutor substitutor = new StringSubstitutor(variablesMap).setVariablePrefix("\"${").setVariableSuffix("}\"");
+        String newQueryContent = substitutor.replace(queryString);
+
+        try {
+            XContentParser parser = XContentType.JSON.xContent()
+                .createParser(queryShardContext.getXContentRegistry(), LoggingDeprecationHandler.INSTANCE, newQueryContent);
+
+            ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
+
+            QueryBuilder newQueryBuilder = parseInnerQueryBuilder(parser);
+
+            return newQueryBuilder;
+
+        } catch (Exception e) {
+            throw new IOException("Failed to rewrite template query: " + newQueryContent, e);
+        }
+    }
+}

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;
@@ -1172,7 +1173,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));
         }

server/src/main/java/org/opensearch/search/pipeline/PipelineProcessingContext.java

@@ -35,4 +35,8 @@ public void setAttribute(String name, Object value) {
     public Object getAttribute(String name) {
         return attributes.get(name);
     }
+
+    public Map<String, Object> getAttributes() {
+        return attributes;
+    }
 }

server/src/test/java/org/opensearch/index/mapper/DateFieldTypeTests.java

@@ -63,6 +63,7 @@
 import org.opensearch.index.mapper.DateFieldMapper.Resolution;
 import org.opensearch.index.mapper.MappedFieldType.Relation;
 import org.opensearch.index.mapper.ParseContext.Document;
+import org.opensearch.index.query.BaseQueryRewriteContext;
 import org.opensearch.index.query.DateRangeIncludingNowQuery;
 import org.opensearch.index.query.QueryRewriteContext;
 import org.opensearch.index.query.QueryShardContext;
@@ -83,7 +84,7 @@ public class DateFieldTypeTests extends FieldTypeTestCase {
     private static final long nowInMillis = 0;
 
     public void testIsFieldWithinRangeEmptyReader() throws IOException {
-        QueryRewriteContext context = new QueryRewriteContext(xContentRegistry(), writableRegistry(), null, () -> nowInMillis);
+        QueryRewriteContext context = new BaseQueryRewriteContext(xContentRegistry(), writableRegistry(), null, () -> nowInMillis);
         IndexReader reader = new MultiReader();
         DateFieldType ft = new DateFieldType("my_date");
         assertEquals(
@@ -120,7 +121,7 @@ public void isFieldWithinRangeTestCase(DateFieldType ft) throws IOException {
         doTestIsFieldWithinQuery(ft, reader, DateTimeZone.UTC, null);
         doTestIsFieldWithinQuery(ft, reader, DateTimeZone.UTC, alternateFormat);
 
-        QueryRewriteContext context = new QueryRewriteContext(xContentRegistry(), writableRegistry(), null, () -> nowInMillis);
+        QueryRewriteContext context = new BaseQueryRewriteContext(xContentRegistry(), writableRegistry(), null, () -> nowInMillis);
 
         // Fields with no value indexed.
         DateFieldType ft2 = new DateFieldType("my_date2");
@@ -132,7 +133,7 @@ public void isFieldWithinRangeTestCase(DateFieldType ft) throws IOException {
 
     private void doTestIsFieldWithinQuery(DateFieldType ft, DirectoryReader reader, DateTimeZone zone, DateMathParser alternateFormat)
         throws IOException {
-        QueryRewriteContext context = new QueryRewriteContext(xContentRegistry(), writableRegistry(), null, () -> nowInMillis);
+        QueryRewriteContext context = new BaseQueryRewriteContext(xContentRegistry(), writableRegistry(), null, () -> nowInMillis);
         assertEquals(
             Relation.INTERSECTS,
             ft.isFieldWithinQuery(reader, "2015-10-09", "2016-01-02", randomBoolean(), randomBoolean(), null, null, context)

server/src/test/java/org/opensearch/index/query/RewriteableTests.java

@@ -45,7 +45,7 @@
 public class RewriteableTests extends OpenSearchTestCase {
 
     public void testRewrite() throws IOException {
-        QueryRewriteContext context = new QueryRewriteContext(null, null, null, null);
+        QueryRewriteContext context = new BaseQueryRewriteContext(null, null, null, null);
         TestRewriteable rewrite = Rewriteable.rewrite(
             new TestRewriteable(randomIntBetween(0, Rewriteable.MAX_REWRITE_ROUNDS)),
             context,
@@ -65,7 +65,7 @@ public void testRewrite() throws IOException {
     }
 
     public void testRewriteAndFetch() throws ExecutionException, InterruptedException {
-        QueryRewriteContext context = new QueryRewriteContext(null, null, null, null);
+        BaseQueryRewriteContext context = new BaseQueryRewriteContext(null, null, null, null);
         PlainActionFuture<TestRewriteable> future = new PlainActionFuture<>();
         Rewriteable.rewriteAndFetch(new TestRewriteable(randomIntBetween(0, Rewriteable.MAX_REWRITE_ROUNDS), true), context, future);
         TestRewriteable rewrite = future.get();
@@ -83,7 +83,7 @@ public void testRewriteAndFetch() throws ExecutionException, InterruptedExceptio
     }
 
     public void testRewriteList() throws IOException {
-        QueryRewriteContext context = new QueryRewriteContext(null, null, null, null);
+        BaseQueryRewriteContext context = new BaseQueryRewriteContext(null, null, null, null);
         List<TestRewriteable> rewriteableList = new ArrayList<>();
         int numInstances = randomIntBetween(1, 10);
         rewriteableList.add(new TestRewriteable(randomIntBetween(1, Rewriteable.MAX_REWRITE_ROUNDS)));