diff --git a/server/src/main/java/org/elasticsearch/index/query/AbstractQueryBuilder.java b/server/src/main/java/org/elasticsearch/index/query/AbstractQueryBuilder.java index 362b80fa571c5..05262798bac2a 100644 --- a/server/src/main/java/org/elasticsearch/index/query/AbstractQueryBuilder.java +++ b/server/src/main/java/org/elasticsearch/index/query/AbstractQueryBuilder.java @@ -283,7 +283,7 @@ public final QueryBuilder rewrite(QueryRewriteContext queryRewriteContext) throw if (queryRewriteInterceptor != null) { var rewritten = queryRewriteInterceptor.interceptAndRewrite(queryRewriteContext, this); if (rewritten != this) { - return rewritten; + return new InterceptedQueryBuilderWrapper(rewritten); } } diff --git a/server/src/main/java/org/elasticsearch/index/query/InterceptedQueryBuilderWrapper.java b/server/src/main/java/org/elasticsearch/index/query/InterceptedQueryBuilderWrapper.java index cb059bf757d42..526439becea60 100644 --- a/server/src/main/java/org/elasticsearch/index/query/InterceptedQueryBuilderWrapper.java +++ b/server/src/main/java/org/elasticsearch/index/query/InterceptedQueryBuilderWrapper.java @@ -21,6 +21,7 @@ /** * Wrapper for instances of {@link QueryBuilder} that have been intercepted using the {@link QueryRewriteInterceptor} to * break out of the rewrite phase. These instances are unwrapped on serialization. + * TODO can this be moved to package protected? */ public class InterceptedQueryBuilderWrapper implements QueryBuilder { @@ -35,6 +36,16 @@ public InterceptedQueryBuilderWrapper(QueryBuilder queryBuilder) { } } + @Override + public QueryBuilder rewrite(QueryRewriteContext queryRewriteContext) throws IOException { + QueryRewriteContext interceptorRemovedContext = queryRewriteContext.getInterceptorRemovedContext(); + QueryBuilder rewritten = queryBuilder.rewrite(interceptorRemovedContext); + if (rewritten != queryBuilder) { + return new InterceptedQueryBuilderWrapper(rewritten); + } + return this; + } + @Override public String getWriteableName() { return queryBuilder.getWriteableName(); diff --git a/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java b/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java index 04b873f391448..9cbaa3bb8fdcb 100644 --- a/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java +++ b/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java @@ -440,4 +440,28 @@ public QueryRewriteInterceptor getQueryRewriteInterceptor() { return queryRewriteInterceptor; } + public QueryRewriteContext getInterceptorRemovedContext() { + if (queryRewriteInterceptor == null) { + return this; + } + return new QueryRewriteContext( + parserConfiguration, + client, + nowInMillis, + mapperService, + mappingLookup, + runtimeMappings, + indexSettings, + fullyQualifiedIndex, + indexNameMatcher, + writeableRegistry, + valuesSourceRegistry, + allowExpensiveQueries, + scriptService, + resolvedIndices, + pit, + null + ); + } + } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMatchQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMatchQueryRewriteInterceptor.java index f402e5fc712bd..1dba609786b9b 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMatchQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMatchQueryRewriteInterceptor.java @@ -12,13 +12,7 @@ import org.elasticsearch.cluster.metadata.InferenceFieldMetadata; import org.elasticsearch.features.NodeFeature; import org.elasticsearch.index.mapper.IndexFieldMapper; -import org.elasticsearch.index.query.BoolQueryBuilder; -import org.elasticsearch.index.query.InterceptedQueryBuilderWrapper; -import org.elasticsearch.index.query.MatchQueryBuilder; -import org.elasticsearch.index.query.QueryBuilder; -import org.elasticsearch.index.query.QueryRewriteContext; -import org.elasticsearch.index.query.TermQueryBuilder; -import org.elasticsearch.index.query.TermsQueryBuilder; +import org.elasticsearch.index.query.*; import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor; import java.util.ArrayList; @@ -64,7 +58,7 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde createSemanticSubQuery(inferenceIndexName, matchQueryBuilder.fieldName(), (String) matchQueryBuilder.value()) ); } - boolQueryBuilder.should(createMatchSubQuery(nonInferenceIndices, matchQueryBuilder.fieldName(), matchQueryBuilder.value())); + boolQueryBuilder.should(createMatchSubQuery(nonInferenceIndices, matchQueryBuilder)); rewritten = boolQueryBuilder; } else { rewritten = new SemanticQueryBuilder(matchQueryBuilder.fieldName(), (String) matchQueryBuilder.value(), false); @@ -87,9 +81,9 @@ private QueryBuilder createSemanticSubQuery(String indexName, String fieldName, return boolQueryBuilder; } - private QueryBuilder createMatchSubQuery(List indices, String fieldName, Object value) { + private QueryBuilder createMatchSubQuery(List indices, MatchQueryBuilder matchQueryBuilder) { BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); - boolQueryBuilder.must(new InterceptedQueryBuilderWrapper(new MatchQueryBuilder(fieldName, value))); + boolQueryBuilder.must(new InterceptedQueryBuilderWrapper(matchQueryBuilder)); boolQueryBuilder.filter(new TermsQueryBuilder(IndexFieldMapper.NAME, indices)); return boolQueryBuilder; }