diff --git a/plugins/ingest-attachment/build.gradle b/plugins/ingest-attachment/build.gradle index 4f30ea9ea7e22..008a776f5a419 100644 --- a/plugins/ingest-attachment/build.gradle +++ b/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 diff --git a/plugins/ingest-attachment/licenses/commons-lang3-3.14.0.jar.sha1 b/plugins/ingest-attachment/licenses/commons-lang3-3.14.0.jar.sha1 deleted file mode 100644 index d783e07e40902..0000000000000 --- a/plugins/ingest-attachment/licenses/commons-lang3-3.14.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -1ed471194b02f2c6cb734a0cd6f6f107c673afae \ No newline at end of file diff --git a/plugins/repository-azure/build.gradle b/plugins/repository-azure/build.gradle index d419f6fafeb30..d91b210d44708 100644 --- a/plugins/repository-azure/build.gradle +++ b/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') } diff --git a/plugins/repository-azure/licenses/commons-lang3-3.14.0.jar.sha1 b/plugins/repository-azure/licenses/commons-lang3-3.14.0.jar.sha1 deleted file mode 100644 index d783e07e40902..0000000000000 --- a/plugins/repository-azure/licenses/commons-lang3-3.14.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -1ed471194b02f2c6cb734a0cd6f6f107c673afae \ No newline at end of file diff --git a/plugins/repository-hdfs/licenses/commons-lang3-3.17.0.jar.sha1 b/plugins/repository-hdfs/licenses/commons-lang3-3.17.0.jar.sha1 deleted file mode 100644 index 073922fda1dbe..0000000000000 --- a/plugins/repository-hdfs/licenses/commons-lang3-3.17.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -b17d2136f0460dcc0d2016ceefca8723bdf4ee70 \ No newline at end of file diff --git a/server/build.gradle b/server/build.gradle index 8dd23491ccd69..2f7fecde2eb44 100644 --- a/server/build.gradle +++ b/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') diff --git a/plugins/ingest-attachment/licenses/commons-lang3-LICENSE.txt b/server/licenses/commons-lang3-LICENSE.txt similarity index 100% rename from plugins/ingest-attachment/licenses/commons-lang3-LICENSE.txt rename to server/licenses/commons-lang3-LICENSE.txt diff --git a/plugins/ingest-attachment/licenses/commons-lang3-NOTICE.txt b/server/licenses/commons-lang3-NOTICE.txt similarity index 100% rename from plugins/ingest-attachment/licenses/commons-lang3-NOTICE.txt rename to server/licenses/commons-lang3-NOTICE.txt diff --git a/server/licenses/commons-text-LICENSE.txt b/server/licenses/commons-text-LICENSE.txt new file mode 100644 index 0000000000000..b687324b56d6a --- /dev/null +++ b/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 diff --git a/server/licenses/commons-text-NOTICE.txt b/server/licenses/commons-text-NOTICE.txt new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/server/src/main/java/org/opensearch/index/query/BaseQueryRewriteContext.java b/server/src/main/java/org/opensearch/index/query/BaseQueryRewriteContext.java index 9d6ce82f6de20..32fc446bfe224 100644 --- a/server/src/main/java/org/opensearch/index/query/BaseQueryRewriteContext.java +++ b/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; diff --git a/server/src/main/java/org/opensearch/index/query/QueryBuilders.java b/server/src/main/java/org/opensearch/index/query/QueryBuilders.java index 387d21830aa38..1debba73136b2 100644 --- a/server/src/main/java/org/opensearch/index/query/QueryBuilders.java +++ b/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 content) { + return new TemplateQueryBuilder(content); + } } diff --git a/server/src/main/java/org/opensearch/index/query/QueryCoordinatorContext.java b/server/src/main/java/org/opensearch/index/query/QueryCoordinatorContext.java index a43d11d0ba7e1..4b2d032f4f9e5 100644 --- a/server/src/main/java/org/opensearch/index/query/QueryCoordinatorContext.java +++ b/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 getContextVariables() { + + Map 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; } } diff --git a/server/src/main/java/org/opensearch/index/query/QueryRewriteContext.java b/server/src/main/java/org/opensearch/index/query/QueryRewriteContext.java index 7159a37dbeec8..aec5914066ab5 100644 --- a/server/src/main/java/org/opensearch/index/query/QueryRewriteContext.java +++ b/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. diff --git a/server/src/main/java/org/opensearch/index/query/TemplateQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/TemplateQueryBuilder.java new file mode 100644 index 0000000000000..4355e50b8628d --- /dev/null +++ b/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 { + public static final String NAME = "template"; + public static final String queryName = "template"; + private final Map content; + + /** + * Constructs a new TemplateQueryBuilder with the given content. + * + * @param content The template content as a map. + */ + public TemplateQueryBuilder(Map 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 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 contextVariables = queryCoordinateContext.getContextVariables(); + String queryString; + + try (XContentBuilder builder = XContentFactory.jsonBuilder()) { + builder.map(this.content); + queryString = builder.toString(); + } + + // Convert Map to Map with proper JSON escaping + Map 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); + } + } +} diff --git a/server/src/main/java/org/opensearch/search/SearchModule.java b/server/src/main/java/org/opensearch/search/SearchModule.java index 3a746259af3b5..facc16cf8c6cd 100644 --- a/server/src/main/java/org/opensearch/search/SearchModule.java +++ b/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 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)); } diff --git a/server/src/main/java/org/opensearch/search/pipeline/PipelineProcessingContext.java b/server/src/main/java/org/opensearch/search/pipeline/PipelineProcessingContext.java index a1f2b8b99d958..ae57692b5a6d0 100644 --- a/server/src/main/java/org/opensearch/search/pipeline/PipelineProcessingContext.java +++ b/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 getAttributes() { + return attributes; + } } diff --git a/server/src/test/java/org/opensearch/index/mapper/DateFieldTypeTests.java b/server/src/test/java/org/opensearch/index/mapper/DateFieldTypeTests.java index 15b16f4610062..7e7e331a29618 100644 --- a/server/src/test/java/org/opensearch/index/mapper/DateFieldTypeTests.java +++ b/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) diff --git a/server/src/test/java/org/opensearch/index/query/RewriteableTests.java b/server/src/test/java/org/opensearch/index/query/RewriteableTests.java index 6385a57f9f370..6e58023ecc7e2 100644 --- a/server/src/test/java/org/opensearch/index/query/RewriteableTests.java +++ b/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 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 rewriteableList = new ArrayList<>(); int numInstances = randomIntBetween(1, 10); rewriteableList.add(new TestRewriteable(randomIntBetween(1, Rewriteable.MAX_REWRITE_ROUNDS))); diff --git a/server/src/test/java/org/opensearch/index/query/TemplateQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/TemplateQueryBuilderTests.java new file mode 100644 index 0000000000000..e3e5d70cd4dcc --- /dev/null +++ b/server/src/test/java/org/opensearch/index/query/TemplateQueryBuilderTests.java @@ -0,0 +1,683 @@ +/* + * 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.geo.GeoPoint; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.common.xcontent.XContentType; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.xcontent.DeprecationHandler; +import org.opensearch.core.xcontent.NamedXContentRegistry; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.search.SearchModule; +import org.opensearch.search.builder.SearchSourceBuilder; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.function.BiConsumer; + +import static org.opensearch.index.query.TemplateQueryBuilder.NAME; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class TemplateQueryBuilderTests extends OpenSearchTestCase { + + /** + * Tests the fromXContent method of TemplateQueryBuilder. + * Verifies that a TemplateQueryBuilder can be correctly created from XContent. + */ + public void testFromXContent() throws IOException { + /* + { + "template": { + "term": { + "message": { + "value": "foo" + } + } + } + } + */ + Map template = new HashMap<>(); + Map term = new HashMap<>(); + Map message = new HashMap<>(); + + message.put("value", "foo"); + term.put("message", message); + template.put("term", term); + + XContentBuilder xContentBuilder = XContentFactory.jsonBuilder().map(template); + + XContentParser contentParser = createParser(xContentBuilder); + contentParser.nextToken(); + TemplateQueryBuilder templateQueryBuilder = TemplateQueryBuilder.fromXContent(contentParser); + + assertEquals(NAME, templateQueryBuilder.getWriteableName()); + assertEquals(template, templateQueryBuilder.getContent()); + + SearchSourceBuilder source = new SearchSourceBuilder().query(templateQueryBuilder); + assertEquals(source.toString(), "{\"query\":{\"template\":{\"term\":{\"message\":{\"value\":\"foo\"}}}}}"); + } + + /** + * Tests the query source generation of TemplateQueryBuilder. + * Verifies that the correct query source is generated from a TemplateQueryBuilder. + */ + public void testQuerySource() { + + Map template = new HashMap<>(); + Map term = new HashMap<>(); + Map message = new HashMap<>(); + + message.put("value", "foo"); + term.put("message", message); + template.put("term", term); + QueryBuilder incomingQuery = new TemplateQueryBuilder(template); + SearchSourceBuilder source = new SearchSourceBuilder().query(incomingQuery); + assertEquals(source.toString(), "{\"query\":{\"template\":{\"term\":{\"message\":{\"value\":\"foo\"}}}}}"); + } + + /** + * Tests parsing a TemplateQueryBuilder from a JSON string. + * Verifies that the parsed query matches the expected structure and can be serialized and deserialized. + */ + public void testFromJson() throws IOException { + String jsonString = "{\n" + + " \"geo_shape\": {\n" + + " \"location\": {\n" + + " \"shape\": {\n" + + " \"type\": \"Envelope\",\n" + + " \"coordinates\": \"${modelPredictionOutcome}\"\n" + + " },\n" + + " \"relation\": \"intersects\"\n" + + " },\n" + + " \"ignore_unmapped\": false,\n" + + " \"boost\": 42.0\n" + + " }\n" + + "}"; + + XContentParser parser = XContentType.JSON.xContent() + .createParser(xContentRegistry(), DeprecationHandler.THROW_UNSUPPORTED_OPERATION, jsonString); + parser.nextToken(); + TemplateQueryBuilder parsed = TemplateQueryBuilder.fromXContent(parser); + + // Check if the parsed query is an instance of TemplateQueryBuilder + assertNotNull(parsed); + assertTrue(parsed instanceof TemplateQueryBuilder); + + // Check if the content of the parsed query matches the expected content + Map expectedContent = new HashMap<>(); + Map geoShape = new HashMap<>(); + Map location = new HashMap<>(); + Map shape = new HashMap<>(); + + shape.put("type", "Envelope"); + shape.put("coordinates", "${modelPredictionOutcome}"); + location.put("shape", shape); + location.put("relation", "intersects"); + geoShape.put("location", location); + geoShape.put("ignore_unmapped", false); + geoShape.put("boost", 42.0); + expectedContent.put("geo_shape", geoShape); + + Map actualContent = new HashMap<>(); + actualContent.put("template", expectedContent); + assertEquals(expectedContent, parsed.getContent()); + + // Test that the query can be serialized and deserialized + BytesStreamOutput out = new BytesStreamOutput(); + parsed.writeTo(out); + StreamInput in = out.bytes().streamInput(); + TemplateQueryBuilder deserializedQuery = new TemplateQueryBuilder(in); + assertEquals(parsed.getContent(), deserializedQuery.getContent()); + + // Test that the query can be converted to XContent + XContentBuilder builder = XContentFactory.jsonBuilder(); + builder.startObject(); + parsed.doXContent(builder, ToXContent.EMPTY_PARAMS); + builder.endObject(); + + Map expectedJson = new HashMap<>(); + Map template = new HashMap<>(); + template.put("geo_shape", geoShape); + expectedJson.put("template", template); + + XContentParser jsonParser = XContentType.JSON.xContent() + .createParser(xContentRegistry(), DeprecationHandler.THROW_UNSUPPORTED_OPERATION, builder.toString()); + Map actualJson = jsonParser.map(); + + assertEquals(expectedJson, actualJson); + } + + /** + * Tests the constructor and getter methods of TemplateQueryBuilder. + * Verifies that the content and writeable name are correctly set and retrieved. + */ + public void testConstructorAndGetters() { + Map content = new HashMap<>(); + content.put("key", "value"); + TemplateQueryBuilder builder = new TemplateQueryBuilder(content); + + assertEquals(content, builder.getContent()); + assertEquals(NAME, builder.getWriteableName()); + } + + /** + * Tests the equals and hashCode methods of TemplateQueryBuilder. + * Verifies that two builders with the same content are equal and have the same hash code, + * while builders with different content are not equal and have different hash codes. + */ + public void testEqualsAndHashCode() { + Map content1 = new HashMap<>(); + content1.put("key", "value"); + TemplateQueryBuilder builder1 = new TemplateQueryBuilder(content1); + + Map content2 = new HashMap<>(); + content2.put("key", "value"); + TemplateQueryBuilder builder2 = new TemplateQueryBuilder(content2); + + Map content3 = new HashMap<>(); + content3.put("key", "different_value"); + TemplateQueryBuilder builder3 = new TemplateQueryBuilder(content3); + + assertTrue(builder1.equals(builder2)); + assertTrue(builder1.hashCode() == builder2.hashCode()); + assertFalse(builder1.equals(builder3)); + assertFalse(builder1.hashCode() == builder3.hashCode()); + } + + /** + * Tests the doToQuery method of TemplateQueryBuilder. + * Verifies that calling doToQuery throws an IllegalStateException. + */ + public void testDoToQuery() { + Map content = new HashMap<>(); + content.put("key", "value"); + TemplateQueryBuilder builder = new TemplateQueryBuilder(content); + + QueryShardContext mockContext = mock(QueryShardContext.class); + expectThrows(IllegalStateException.class, () -> builder.doToQuery(mockContext)); + } + + /** + * Tests the serialization and deserialization of TemplateQueryBuilder. + * Verifies that a builder can be written to a stream and read back correctly. + */ + public void testStreamRoundTrip() throws IOException { + Map content = new HashMap<>(); + content.put("key", "value"); + TemplateQueryBuilder original = new TemplateQueryBuilder(content); + + BytesStreamOutput out = new BytesStreamOutput(); + original.writeTo(out); + + StreamInput in = out.bytes().streamInput(); + TemplateQueryBuilder deserialized = new TemplateQueryBuilder(in); + + assertEquals(original, deserialized); + } + + /** + * Tests the doRewrite method of TemplateQueryBuilder with a simple term query. + * Verifies that the template is correctly rewritten to a TermQueryBuilder. + */ + public void testDoRewrite() throws IOException { + + Map template = new HashMap<>(); + Map term = new HashMap<>(); + Map message = new HashMap<>(); + + message.put("value", "foo"); + term.put("message", message); + template.put("term", term); + TemplateQueryBuilder templateQueryBuilder = new TemplateQueryBuilder(template); + TermQueryBuilder termQueryBuilder = new TermQueryBuilder("message", "foo"); + + QueryCoordinatorContext queryRewriteContext = mockQueryRewriteContext(); + + Map contextVariables = new HashMap<>(); + when(queryRewriteContext.getContextVariables()).thenReturn(contextVariables); + + TermQueryBuilder newQuery = (TermQueryBuilder) templateQueryBuilder.doRewrite(queryRewriteContext); + + assertEquals(newQuery, termQueryBuilder); + assertEquals( + "{\n" + + " \"term\" : {\n" + + " \"message\" : {\n" + + " \"value\" : \"foo\",\n" + + " \"boost\" : 1.0\n" + + " }\n" + + " }\n" + + "}", + newQuery.toString() + ); + } + + /** + * Tests the doRewrite method of TemplateQueryBuilder with a string variable. + * Verifies that the template is correctly rewritten with the variable substituted. + */ + public void testDoRewriteWithString() throws IOException { + + Map template = new HashMap<>(); + Map term = new HashMap<>(); + Map message = new HashMap<>(); + + message.put("value", "${response}"); + term.put("message", message); + template.put("term", term); + TemplateQueryBuilder templateQueryBuilder = new TemplateQueryBuilder(template); + TermQueryBuilder termQueryBuilder = new TermQueryBuilder("message", "foo"); + + QueryCoordinatorContext queryRewriteContext = mockQueryRewriteContext(); + + Map contextVariables = new HashMap<>(); + contextVariables.put("response", "foo"); + when(queryRewriteContext.getContextVariables()).thenReturn(contextVariables); + + TermQueryBuilder newQuery = (TermQueryBuilder) templateQueryBuilder.doRewrite(queryRewriteContext); + + assertEquals(newQuery, termQueryBuilder); + assertEquals( + "{\n" + + " \"term\" : {\n" + + " \"message\" : {\n" + + " \"value\" : \"foo\",\n" + + " \"boost\" : 1.0\n" + + " }\n" + + " }\n" + + "}", + newQuery.toString() + ); + } + + /** + * Tests the doRewrite method of TemplateQueryBuilder with a list variable. + * Verifies that the template is correctly rewritten with the list variable substituted. + */ + public void testDoRewriteWithList() throws IOException { + ArrayList termsList = new ArrayList<>(); + termsList.add("foo"); + termsList.add("bar"); + + Map template = new HashMap<>(); + Map terms = new HashMap<>(); + + terms.put("message", "${response}"); + template.put("terms", terms); + TemplateQueryBuilder templateQueryBuilder = new TemplateQueryBuilder(template); + + TermsQueryBuilder termsQueryBuilder = new TermsQueryBuilder("message", termsList); + + QueryCoordinatorContext queryRewriteContext = mockQueryRewriteContext(); + + Map contextVariables = new HashMap<>(); + contextVariables.put("response", termsList); + when(queryRewriteContext.getContextVariables()).thenReturn(contextVariables); + NamedXContentRegistry TEST_XCONTENT_REGISTRY_FOR_QUERY = new NamedXContentRegistry( + new SearchModule(Settings.EMPTY, List.of()).getNamedXContents() + ); + when(queryRewriteContext.getXContentRegistry()).thenReturn(TEST_XCONTENT_REGISTRY_FOR_QUERY); + TermsQueryBuilder newQuery = (TermsQueryBuilder) templateQueryBuilder.doRewrite(queryRewriteContext); + assertEquals(newQuery, termsQueryBuilder); + assertEquals( + "{\n" + + " \"terms\" : {\n" + + " \"message\" : [\n" + + " \"foo\",\n" + + " \"bar\"\n" + + " ],\n" + + " \"boost\" : 1.0\n" + + " }\n" + + "}", + newQuery.toString() + ); + } + + /** + * Tests the doRewrite method of TemplateQueryBuilder with a geo_distance query. + * Verifies that the template is correctly rewritten for a geo_distance query. + */ + public void testDoRewriteWithGeoDistanceQuery() throws IOException { + Map template = new HashMap<>(); + Map geoDistance = new HashMap<>(); + + geoDistance.put("distance", "12km"); + geoDistance.put("pin.location", "${geoPoint}"); + template.put("geo_distance", geoDistance); + + TemplateQueryBuilder templateQueryBuilder = new TemplateQueryBuilder(template); + + GeoPoint geoPoint = new GeoPoint(40, -70); + + QueryCoordinatorContext queryRewriteContext = mockQueryRewriteContext(); + Map contextVariables = new HashMap<>(); + contextVariables.put("geoPoint", geoPoint); + when(queryRewriteContext.getContextVariables()).thenReturn(contextVariables); + + GeoDistanceQueryBuilder expectedQuery = new GeoDistanceQueryBuilder("pin.location"); + expectedQuery.point(geoPoint).distance("12km"); + + QueryBuilder newQuery = templateQueryBuilder.doRewrite(queryRewriteContext); + assertEquals(expectedQuery, newQuery); + assertEquals( + "{\n" + + " \"geo_distance\" : {\n" + + " \"pin.location\" : [\n" + + " -70.0,\n" + + " 40.0\n" + + " ],\n" + + " \"distance\" : 12000.0,\n" + + " \"distance_type\" : \"arc\",\n" + + " \"validation_method\" : \"STRICT\",\n" + + " \"ignore_unmapped\" : false,\n" + + " \"boost\" : 1.0\n" + + " }\n" + + "}", + newQuery.toString() + ); + } + + /** + * Tests the doRewrite method of TemplateQueryBuilder with a range query. + * Verifies that the template is correctly rewritten for a range query. + */ + public void testDoRewriteWithRangeQuery() throws IOException { + Map template = new HashMap<>(); + Map range = new HashMap<>(); + Map age = new HashMap<>(); + + age.put("gte", "${minAge}"); + age.put("lte", "${maxAge}"); + range.put("age", age); + template.put("range", range); + + TemplateQueryBuilder templateQueryBuilder = new TemplateQueryBuilder(template); + + QueryCoordinatorContext queryRewriteContext = mockQueryRewriteContext(); + Map contextVariables = new HashMap<>(); + contextVariables.put("minAge", 25); + contextVariables.put("maxAge", 35); + when(queryRewriteContext.getContextVariables()).thenReturn(contextVariables); + + RangeQueryBuilder expectedQuery = new RangeQueryBuilder("age"); + expectedQuery.gte(25).lte(35); + + QueryBuilder newQuery = templateQueryBuilder.doRewrite(queryRewriteContext); + assertEquals(expectedQuery, newQuery); + assertEquals( + "{\n" + + " \"range\" : {\n" + + " \"age\" : {\n" + + " \"from\" : 25,\n" + + " \"to\" : 35,\n" + + " \"include_lower\" : true,\n" + + " \"include_upper\" : true,\n" + + " \"boost\" : 1.0\n" + + " }\n" + + " }\n" + + "}", + newQuery.toString() + ); + } + + /** + * Tests the doRewrite method of TemplateQueryBuilder with a nested map variable. + * Verifies that the template is correctly rewritten with the nested map variable substituted. + */ + public void testDoRewriteWithNestedMap() throws IOException { + Map template = new HashMap<>(); + Map bool = new HashMap<>(); + List> must = new ArrayList<>(); + Map match = new HashMap<>(); + Map textEntry = new HashMap<>(); + + textEntry.put("text_entry", "${keyword}"); + match.put("match", textEntry); + must.add(match); + bool.put("must", must); + + List> should = new ArrayList<>(); + Map shouldMatch1 = new HashMap<>(); + Map shouldTextEntry1 = new HashMap<>(); + shouldTextEntry1.put("text_entry", "life"); + shouldMatch1.put("match", shouldTextEntry1); + should.add(shouldMatch1); + + Map shouldMatch2 = new HashMap<>(); + Map shouldTextEntry2 = new HashMap<>(); + shouldTextEntry2.put("text_entry", "grace"); + shouldMatch2.put("match", shouldTextEntry2); + should.add(shouldMatch2); + + bool.put("should", should); + bool.put("minimum_should_match", 1); + + Map filter = new HashMap<>(); + Map term = new HashMap<>(); + term.put("play_name", "Romeo and Juliet"); + filter.put("term", term); + bool.put("filter", filter); + + template.put("bool", bool); + + TemplateQueryBuilder templateQueryBuilder = new TemplateQueryBuilder(template); + + QueryCoordinatorContext queryRewriteContext = mockQueryRewriteContext(); + Map contextVariables = new HashMap<>(); + contextVariables.put("keyword", "love"); + when(queryRewriteContext.getContextVariables()).thenReturn(contextVariables); + + BoolQueryBuilder expectedQuery = new BoolQueryBuilder().must(new MatchQueryBuilder("text_entry", "love")) + .should(new MatchQueryBuilder("text_entry", "life")) + .should(new MatchQueryBuilder("text_entry", "grace")) + .filter(new TermQueryBuilder("play_name", "Romeo and Juliet")) + .minimumShouldMatch(1); + + QueryBuilder newQuery = templateQueryBuilder.doRewrite(queryRewriteContext); + assertEquals(expectedQuery, newQuery); + assertEquals( + "{\n" + + " \"bool\" : {\n" + + " \"must\" : [\n" + + " {\n" + + " \"match\" : {\n" + + " \"text_entry\" : {\n" + + " \"query\" : \"love\",\n" + + " \"operator\" : \"OR\",\n" + + " \"prefix_length\" : 0,\n" + + " \"max_expansions\" : 50,\n" + + " \"fuzzy_transpositions\" : true,\n" + + " \"lenient\" : false,\n" + + " \"zero_terms_query\" : \"NONE\",\n" + + " \"auto_generate_synonyms_phrase_query\" : true,\n" + + " \"boost\" : 1.0\n" + + " }\n" + + " }\n" + + " }\n" + + " ],\n" + + " \"filter\" : [\n" + + " {\n" + + " \"term\" : {\n" + + " \"play_name\" : {\n" + + " \"value\" : \"Romeo and Juliet\",\n" + + " \"boost\" : 1.0\n" + + " }\n" + + " }\n" + + " }\n" + + " ],\n" + + " \"should\" : [\n" + + " {\n" + + " \"match\" : {\n" + + " \"text_entry\" : {\n" + + " \"query\" : \"life\",\n" + + " \"operator\" : \"OR\",\n" + + " \"prefix_length\" : 0,\n" + + " \"max_expansions\" : 50,\n" + + " \"fuzzy_transpositions\" : true,\n" + + " \"lenient\" : false,\n" + + " \"zero_terms_query\" : \"NONE\",\n" + + " \"auto_generate_synonyms_phrase_query\" : true,\n" + + " \"boost\" : 1.0\n" + + " }\n" + + " }\n" + + " },\n" + + " {\n" + + " \"match\" : {\n" + + " \"text_entry\" : {\n" + + " \"query\" : \"grace\",\n" + + " \"operator\" : \"OR\",\n" + + " \"prefix_length\" : 0,\n" + + " \"max_expansions\" : 50,\n" + + " \"fuzzy_transpositions\" : true,\n" + + " \"lenient\" : false,\n" + + " \"zero_terms_query\" : \"NONE\",\n" + + " \"auto_generate_synonyms_phrase_query\" : true,\n" + + " \"boost\" : 1.0\n" + + " }\n" + + " }\n" + + " }\n" + + " ],\n" + + " \"adjust_pure_negative\" : true,\n" + + " \"minimum_should_match\" : \"1\",\n" + + " \"boost\" : 1.0\n" + + " }\n" + + "}", + newQuery.toString() + ); + } + + /** + * Tests the doRewrite method with an invalid query type. + * Verifies that an IOException is thrown when an invalid query type is used. + */ + public void testDoRewriteWithInvalidQueryType() throws IOException { + Map template = new HashMap<>(); + template.put("invalid_query_type", new HashMap<>()); + TemplateQueryBuilder templateQueryBuilder = new TemplateQueryBuilder(template); + + QueryCoordinatorContext queryRewriteContext = mockQueryRewriteContext(); + when(queryRewriteContext.getContextVariables()).thenReturn(new HashMap<>()); + + IOException exception = expectThrows(IOException.class, () -> templateQueryBuilder.doRewrite(queryRewriteContext)); + assertTrue(exception.getMessage().contains("Failed to rewrite template query")); + } + + /** + * Tests the doRewrite method with a malformed JSON query. + * Verifies that an IOException is thrown when the query JSON is malformed. + */ + public void testDoRewriteWithMalformedJson() throws IOException { + Map template = new HashMap<>(); + template.put("malformed_json", "{ this is not valid JSON }"); + TemplateQueryBuilder templateQueryBuilder = new TemplateQueryBuilder(template); + + QueryCoordinatorContext queryRewriteContext = mockQueryRewriteContext(); + when(queryRewriteContext.getContextVariables()).thenReturn(new HashMap<>()); + + IOException exception = expectThrows(IOException.class, () -> templateQueryBuilder.doRewrite(queryRewriteContext)); + assertTrue(exception.getMessage().contains("Failed to rewrite template query")); + } + + /** + * Tests the doRewrite method with an invalid matchall query. + * Verifies that an IOException is thrown when an invalid matchall query is used. + */ + public void testDoRewriteWithInvalidMatchAllQuery() throws IOException { + Map template = new HashMap<>(); + template.put("matchall_1", new HashMap<>()); + TemplateQueryBuilder templateQueryBuilder = new TemplateQueryBuilder(template); + + QueryCoordinatorContext queryRewriteContext = mockQueryRewriteContext(); + when(queryRewriteContext.getContextVariables()).thenReturn(new HashMap<>()); + + IOException exception = expectThrows(IOException.class, () -> templateQueryBuilder.doRewrite(queryRewriteContext)); + assertTrue(exception.getMessage().contains("Failed to rewrite template query")); + } + + /** + * Tests the doRewrite method with a missing required field in a query. + * Verifies that an IOException is thrown when a required field is missing. + */ + public void testDoRewriteWithMissingRequiredField() throws IOException { + Map template = new HashMap<>(); + template.put("term", "value");// Missing the required field for term query + TemplateQueryBuilder templateQueryBuilder = new TemplateQueryBuilder(template); + + QueryCoordinatorContext queryRewriteContext = mockQueryRewriteContext(); + when(queryRewriteContext.getContextVariables()).thenReturn(new HashMap<>()); + + IOException exception = expectThrows(IOException.class, () -> templateQueryBuilder.doRewrite(queryRewriteContext)); + assertTrue(exception.getMessage().contains("Failed to rewrite template query")); + } + + /** + * Tests the doRewrite method with a malformed variable substitution. + * Verifies that an IOException is thrown when a malformed variable is used. + */ + public void testDoRewriteWithMalformedVariableSubstitution() throws IOException { + + Map template = new HashMap<>(); + Map terms = new HashMap<>(); + + terms.put("message", "${malformed_variable}"); + template.put("terms", terms); + TemplateQueryBuilder templateQueryBuilder = new TemplateQueryBuilder(template); + + QueryCoordinatorContext queryRewriteContext = mockQueryRewriteContext(); + + Map contextVariables = new HashMap<>(); + contextVariables.put("response", "should be a list but this is a string"); + + when(queryRewriteContext.getContextVariables()).thenReturn(contextVariables); + + IOException exception = expectThrows(IOException.class, () -> templateQueryBuilder.doRewrite(queryRewriteContext)); + assertTrue(exception.getMessage().contains("Failed to rewrite template query")); + } + + /** + * Helper method to create a mock QueryCoordinatorContext for testing. + */ + private QueryCoordinatorContext mockQueryRewriteContext() { + QueryCoordinatorContext queryRewriteContext = mock(QueryCoordinatorContext.class); + final CountDownLatch inProgressLatch = new CountDownLatch(1); + doAnswer(invocation -> { + BiConsumer> biConsumer = invocation.getArgument(0); + biConsumer.accept( + null, + ActionListener.wrap( + response -> inProgressLatch.countDown(), + err -> fail("Failed to set query tokens supplier: " + err.getMessage()) + ) + ); + return null; + }).when(queryRewriteContext).registerAsyncAction(any()); + + NamedXContentRegistry TEST_XCONTENT_REGISTRY_FOR_QUERY = new NamedXContentRegistry( + new SearchModule(Settings.EMPTY, List.of()).getNamedXContents() + ); + when(queryRewriteContext.getXContentRegistry()).thenReturn(TEST_XCONTENT_REGISTRY_FOR_QUERY); + + return queryRewriteContext; + } +} diff --git a/server/src/test/java/org/opensearch/search/SearchModuleTests.java b/server/src/test/java/org/opensearch/search/SearchModuleTests.java index 9e8e7afe332f1..d78393e917b2f 100644 --- a/server/src/test/java/org/opensearch/search/SearchModuleTests.java +++ b/server/src/test/java/org/opensearch/search/SearchModuleTests.java @@ -608,7 +608,8 @@ public Optional create(IndexSettings indexSettin "terms_set", "wildcard", "wrapper", - "distance_feature" }; + "distance_feature", + "template" }; // add here deprecated queries to make sure we log a deprecation warnings when they are used private static final String[] DEPRECATED_QUERIES = new String[] { "common", "field_masking_span" }; diff --git a/server/src/test/java/org/opensearch/search/aggregations/AggregatorFactoriesTests.java b/server/src/test/java/org/opensearch/search/aggregations/AggregatorFactoriesTests.java index c930d27b068f8..a5724d3c34352 100644 --- a/server/src/test/java/org/opensearch/search/aggregations/AggregatorFactoriesTests.java +++ b/server/src/test/java/org/opensearch/search/aggregations/AggregatorFactoriesTests.java @@ -45,6 +45,7 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.env.Environment; +import org.opensearch.index.query.BaseQueryRewriteContext; import org.opensearch.index.query.MatchAllQueryBuilder; import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.QueryRewriteContext; @@ -255,7 +256,7 @@ public void testRewriteAggregation() throws Exception { BucketScriptPipelineAggregationBuilder pipelineAgg = new BucketScriptPipelineAggregationBuilder("const", new Script("1")); AggregatorFactories.Builder builder = new AggregatorFactories.Builder().addAggregator(filterAggBuilder) .addPipelineAggregator(pipelineAgg); - AggregatorFactories.Builder rewritten = builder.rewrite(new QueryRewriteContext(xContentRegistry, null, null, () -> 0L)); + AggregatorFactories.Builder rewritten = builder.rewrite(new BaseQueryRewriteContext(xContentRegistry, null, null, () -> 0L)); assertNotSame(builder, rewritten); Collection aggregatorFactories = rewritten.getAggregatorFactories(); assertEquals(1, aggregatorFactories.size()); @@ -268,7 +269,9 @@ public void testRewriteAggregation() throws Exception { assertThat(rewrittenFilter, instanceOf(TermsQueryBuilder.class)); // Check that a further rewrite returns the same aggregation factories builder - AggregatorFactories.Builder secondRewritten = rewritten.rewrite(new QueryRewriteContext(xContentRegistry, null, null, () -> 0L)); + AggregatorFactories.Builder secondRewritten = rewritten.rewrite( + new BaseQueryRewriteContext(xContentRegistry, null, null, () -> 0L) + ); assertSame(rewritten, secondRewritten); } @@ -277,7 +280,7 @@ public void testRewritePipelineAggregationUnderAggregation() throws Exception { new RewrittenPipelineAggregationBuilder() ); AggregatorFactories.Builder builder = new AggregatorFactories.Builder().addAggregator(filterAggBuilder); - QueryRewriteContext context = new QueryRewriteContext(xContentRegistry, null, null, () -> 0L); + QueryRewriteContext context = new BaseQueryRewriteContext(xContentRegistry, null, null, () -> 0L); AggregatorFactories.Builder rewritten = builder.rewrite(context); CountDownLatch latch = new CountDownLatch(1); context.executeAsyncActions(new ActionListener() { @@ -304,7 +307,7 @@ public void testRewriteAggregationAtTopLevel() throws Exception { FilterAggregationBuilder filterAggBuilder = new FilterAggregationBuilder("titles", new MatchAllQueryBuilder()); AggregatorFactories.Builder builder = new AggregatorFactories.Builder().addAggregator(filterAggBuilder) .addPipelineAggregator(new RewrittenPipelineAggregationBuilder()); - QueryRewriteContext context = new QueryRewriteContext(xContentRegistry, null, null, () -> 0L); + QueryRewriteContext context = new BaseQueryRewriteContext(xContentRegistry, null, null, () -> 0L); AggregatorFactories.Builder rewritten = builder.rewrite(context); CountDownLatch latch = new CountDownLatch(1); context.executeAsyncActions(new ActionListener() { diff --git a/server/src/test/java/org/opensearch/search/aggregations/bucket/FiltersTests.java b/server/src/test/java/org/opensearch/search/aggregations/bucket/FiltersTests.java index 56f7f450dbdfb..770f18f781689 100644 --- a/server/src/test/java/org/opensearch/search/aggregations/bucket/FiltersTests.java +++ b/server/src/test/java/org/opensearch/search/aggregations/bucket/FiltersTests.java @@ -36,12 +36,12 @@ import org.opensearch.core.xcontent.MediaTypeRegistry; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.index.query.BaseQueryRewriteContext; import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.MatchAllQueryBuilder; import org.opensearch.index.query.MatchNoneQueryBuilder; import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.QueryBuilders; -import org.opensearch.index.query.QueryRewriteContext; import org.opensearch.search.aggregations.AggregationBuilder; import org.opensearch.search.aggregations.BaseAggregationTestCase; import org.opensearch.search.aggregations.bucket.filter.FiltersAggregationBuilder; @@ -147,12 +147,12 @@ public void testRewrite() throws IOException { // test non-keyed filter that doesn't rewrite AggregationBuilder original = new FiltersAggregationBuilder("my-agg", new MatchAllQueryBuilder()); original.setMetadata(Collections.singletonMap(randomAlphaOfLengthBetween(1, 20), randomAlphaOfLengthBetween(1, 20))); - AggregationBuilder rewritten = original.rewrite(new QueryRewriteContext(xContentRegistry(), null, null, () -> 0L)); + AggregationBuilder rewritten = original.rewrite(new BaseQueryRewriteContext(xContentRegistry(), null, null, () -> 0L)); assertSame(original, rewritten); // test non-keyed filter that does rewrite original = new FiltersAggregationBuilder("my-agg", new BoolQueryBuilder()); - rewritten = original.rewrite(new QueryRewriteContext(xContentRegistry(), null, null, () -> 0L)); + rewritten = original.rewrite(new BaseQueryRewriteContext(xContentRegistry(), null, null, () -> 0L)); assertNotSame(original, rewritten); assertThat(rewritten, instanceOf(FiltersAggregationBuilder.class)); assertEquals("my-agg", ((FiltersAggregationBuilder) rewritten).getName()); @@ -163,12 +163,12 @@ public void testRewrite() throws IOException { // test keyed filter that doesn't rewrite original = new FiltersAggregationBuilder("my-agg", new KeyedFilter("my-filter", new MatchAllQueryBuilder())); - rewritten = original.rewrite(new QueryRewriteContext(xContentRegistry(), null, null, () -> 0L)); + rewritten = original.rewrite(new BaseQueryRewriteContext(xContentRegistry(), null, null, () -> 0L)); assertSame(original, rewritten); // test non-keyed filter that does rewrite original = new FiltersAggregationBuilder("my-agg", new KeyedFilter("my-filter", new BoolQueryBuilder())); - rewritten = original.rewrite(new QueryRewriteContext(xContentRegistry(), null, null, () -> 0L)); + rewritten = original.rewrite(new BaseQueryRewriteContext(xContentRegistry(), null, null, () -> 0L)); assertNotSame(original, rewritten); assertThat(rewritten, instanceOf(FiltersAggregationBuilder.class)); assertEquals("my-agg", ((FiltersAggregationBuilder) rewritten).getName()); @@ -180,7 +180,7 @@ public void testRewrite() throws IOException { // test sub-agg filter that does rewrite original = new TermsAggregationBuilder("terms").userValueTypeHint(ValueType.BOOLEAN) .subAggregation(new FiltersAggregationBuilder("my-agg", new KeyedFilter("my-filter", new BoolQueryBuilder()))); - rewritten = original.rewrite(new QueryRewriteContext(xContentRegistry(), null, null, () -> 0L)); + rewritten = original.rewrite(new BaseQueryRewriteContext(xContentRegistry(), null, null, () -> 0L)); assertNotSame(original, rewritten); assertNotEquals(original, rewritten); assertThat(rewritten, instanceOf(TermsAggregationBuilder.class)); @@ -189,7 +189,7 @@ public void testRewrite() throws IOException { assertThat(subAgg, instanceOf(FiltersAggregationBuilder.class)); assertNotSame(original.getSubAggregations().iterator().next(), subAgg); assertEquals("my-agg", subAgg.getName()); - assertSame(rewritten, rewritten.rewrite(new QueryRewriteContext(xContentRegistry(), null, null, () -> 0L))); + assertSame(rewritten, rewritten.rewrite(new BaseQueryRewriteContext(xContentRegistry(), null, null, () -> 0L))); } public void testRewritePreservesOtherBucket() throws IOException { @@ -197,7 +197,7 @@ public void testRewritePreservesOtherBucket() throws IOException { originalFilters.otherBucket(randomBoolean()); originalFilters.otherBucketKey(randomAlphaOfLength(10)); - AggregationBuilder rewritten = originalFilters.rewrite(new QueryRewriteContext(xContentRegistry(), null, null, () -> 0L)); + AggregationBuilder rewritten = originalFilters.rewrite(new BaseQueryRewriteContext(xContentRegistry(), null, null, () -> 0L)); assertThat(rewritten, instanceOf(FiltersAggregationBuilder.class)); FiltersAggregationBuilder rewrittenFilters = (FiltersAggregationBuilder) rewritten; diff --git a/server/src/test/java/org/opensearch/search/builder/SearchSourceBuilderTests.java b/server/src/test/java/org/opensearch/search/builder/SearchSourceBuilderTests.java index da8ccc9e121e0..025c08ecd703b 100644 --- a/server/src/test/java/org/opensearch/search/builder/SearchSourceBuilderTests.java +++ b/server/src/test/java/org/opensearch/search/builder/SearchSourceBuilderTests.java @@ -47,10 +47,10 @@ import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.index.query.BaseQueryRewriteContext; import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.MatchNoneQueryBuilder; import org.opensearch.index.query.QueryBuilders; -import org.opensearch.index.query.QueryRewriteContext; import org.opensearch.index.query.RandomQueryBuilder; import org.opensearch.index.query.Rewriteable; import org.opensearch.script.Script; @@ -713,7 +713,7 @@ private void assertIndicesBoostParseErrorMessage(String restContent, String expe private SearchSourceBuilder rewrite(SearchSourceBuilder searchSourceBuilder) throws IOException { return Rewriteable.rewrite( searchSourceBuilder, - new QueryRewriteContext(xContentRegistry(), writableRegistry(), null, Long.valueOf(1)::longValue) + new BaseQueryRewriteContext(xContentRegistry(), writableRegistry(), null, Long.valueOf(1)::longValue) ); } }