add visitor logic to FunctionScoreQueryBuilder

Signed-off-by: jdnvn <[email protected]>
opensearch-project · Dec 4, 2024 · 10034b5 · 10034b5
1 parent b1bf72f
commit 10034b5
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/server/src/main/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilder.java
@@ -32,6 +32,7 @@
 
 package org.opensearch.index.query.functionscore;
 
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.Query;
 import org.opensearch.common.Nullable;
@@ -52,6 +53,7 @@
 import org.opensearch.index.query.MatchAllQueryBuilder;
 import org.opensearch.index.query.MatchNoneQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.QueryBuilderVisitor;
 import org.opensearch.index.query.QueryRewriteContext;
 import org.opensearch.index.query.QueryShardContext;
 
@@ -704,4 +706,12 @@ private static String parseFiltersAndFunctions(
         }
         return currentFieldName;
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        if (query != null) {
+            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(query);
+        }
+    }
 }
diff --git a/...rc/test/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilderTests.java b/...rc/test/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilderTests.java
@@ -75,6 +75,7 @@
 import org.hamcrest.Matcher;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
@@ -938,4 +939,14 @@ public void testMustRewrite() throws IOException {
         e = expectThrows(IllegalStateException.class, () -> functionQueryBuilder2.toQuery(context));
         assertEquals("Rewrite first", e.getMessage());
     }
+
+    public void testVisit() {
+        TermQueryBuilder termQueryBuilder = new TermQueryBuilder("unmapped_field", "foo");
+        FunctionScoreQueryBuilder builder = new FunctionScoreQueryBuilder(termQueryBuilder);
+
+        List<QueryBuilder> visitedQueries = new ArrayList<>();
+        builder.visit(createTestVisitor(visitedQueries));
+
+        assertEquals(2, visitedQueries.size());
+    }
 }