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 415a1ca2eb083..89e8e1f9adbdb 100644 --- a/server/src/main/java/org/elasticsearch/index/query/InterceptedQueryBuilderWrapper.java +++ b/server/src/main/java/org/elasticsearch/index/query/InterceptedQueryBuilderWrapper.java @@ -45,6 +45,17 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException { return queryBuilder.doToQuery(context); } + @Override + protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws IOException { + // Ensure that we only perform interception once in a query's rewrite phase + QueryRewriteContext interceptedContext = queryRewriteContext.getInterceptedQueryRewriteContext(); + var rewritten = queryBuilder.rewrite(interceptedContext); + if (rewritten == queryBuilder) { + return this; + } + return rewritten; + } + @Override protected boolean doEquals(T other) { // Handle the edge case where we need to unwrap the incoming query builder 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 e3367dd3e8c9f..26d284a647d67 100644 --- a/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java +++ b/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java @@ -159,6 +159,32 @@ public QueryRewriteContext( ); } + public QueryRewriteContext getInterceptedQueryRewriteContext() { + + if (queryRewriteInterceptor == null) { + return this; + } + + return new QueryRewriteContext( + parserConfiguration, + client, + nowInMillis, + mapperService, + mappingLookup, + runtimeMappings, + indexSettings, + fullyQualifiedIndex, + indexNameMatcher, + writeableRegistry, + valuesSourceRegistry, + allowExpensiveQueries, + scriptService, + resolvedIndices, + pit, + null + ); + } + /** * The registry used to build new {@link XContentParser}s. Contains registered named parsers needed to parse the query. * @@ -438,4 +464,5 @@ public String getTierPreference() { public QueryRewriteInterceptor getQueryRewriteInterceptor() { return queryRewriteInterceptor; } + }