Script Score Query Builder Visit Method (opensearch-project#10198)

* Script Score Builder Visit Method Signed-off-by: Varun Jain <varunudr@amazon.com> * ScriptScore Query Builder Visit Method Signed-off-by: Varun Jain <varunudr@amazon.com> * gradlew spotlessApply Signed-off-by: Varun Jain <varunudr@amazon.com> --------- Signed-off-by: Varun Jain <varunudr@amazon.com>
brusic · Sep 25, 2023 · a8969cb · a8969cb
1 parent f6646fb
commit a8969cb
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/server/src/main/java/org/opensearch/index/query/functionscore/ScriptScoreQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/functionscore/ScriptScoreQueryBuilder.java
@@ -32,6 +32,7 @@
 
 package org.opensearch.index.query.functionscore;
 
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.Query;
 import org.opensearch.OpenSearchException;
 import org.opensearch.common.lucene.search.function.ScriptScoreQuery;
@@ -45,6 +46,7 @@
 import org.opensearch.index.query.InnerHitContextBuilder;
 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;
 import org.opensearch.script.ScoreScript;
@@ -224,4 +226,12 @@ protected void extractInnerHitBuilders(Map<String, InnerHitContextBuilder> inner
         InnerHitContextBuilder.extractInnerHits(query(), innerHits);
     }
 
+    @Override
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        if (query != null) {
+            QueryBuilderVisitor subVisitor = visitor.getChildVisitor(BooleanClause.Occur.MUST);
+            subVisitor.accept(query);
+        }
+    }
 }
diff --git a/server/src/test/java/org/opensearch/index/query/ScriptScoreQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/ScriptScoreQueryBuilderTests.java
@@ -43,7 +43,9 @@
 import org.opensearch.test.AbstractQueryTestCase;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.List;
 
 import static org.opensearch.index.query.QueryBuilders.matchAllQuery;
 import static org.hamcrest.CoreMatchers.instanceOf;
@@ -140,4 +142,11 @@ public void testDisallowExpensiveQueries() {
         OpenSearchException e = expectThrows(OpenSearchException.class, () -> queryBuilder.toQuery(queryShardContext));
         assertEquals("[script score] queries cannot be executed when 'search.allow_expensive_queries' is set to false.", e.getMessage());
     }
+
+    public void testVisit() {
+        ScriptScoreQueryBuilder scriptQueryBuilder = doCreateTestQueryBuilder();
+        List<QueryBuilder> visitedQueries = new ArrayList<>();
+        scriptQueryBuilder.visit(createTestVisitor(visitedQueries));
+        assertEquals(2, visitedQueries.size());
+    }
 }