DOESN'T WORK - for testing

kderusso · Dec 10, 2024 · 2830b2d · 2830b2d
1 parent af4cc21
commit 2830b2d
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 11 deletions.
diff --git 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
@@ -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
@@ -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/.../java/org/elasticsearch/xpack/inference/queries/SemanticMatchQueryRewriteInterceptor.java b/.../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<String> indices, String fieldName, Object value) {
+    private QueryBuilder createMatchSubQuery(List<String> 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;
     }