Ensure interception happens only once

kderusso · Dec 6, 2024 · b3b2507 · b3b2507
1 parent 54989ab
commit b3b2507
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 0 deletions.
diff --git 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
@@ -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;
     }
+
 }