From 0dacef7d9aa0e785a7792180fd939c23fc730ccf Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Mon, 25 Sep 2023 23:29:36 +0000 Subject: [PATCH] Script Score Query Builder Visit Method (#10198) * Script Score Builder Visit Method Signed-off-by: Varun Jain * ScriptScore Query Builder Visit Method Signed-off-by: Varun Jain * gradlew spotlessApply Signed-off-by: Varun Jain --------- Signed-off-by: Varun Jain (cherry picked from commit a8969cb11f3a03a9d7f982cb1fab0525ca957535) Signed-off-by: github-actions[bot] --- .../query/functionscore/ScriptScoreQueryBuilder.java | 10 ++++++++++ .../index/query/ScriptScoreQueryBuilderTests.java | 9 +++++++++ 2 files changed, 19 insertions(+) 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 index a7e2360325c5a..9d28d81f732df 100644 --- 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.LegacyESVersion; import org.opensearch.OpenSearchException; @@ -46,6 +47,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; @@ -233,4 +235,12 @@ protected void extractInnerHitBuilders(Map 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 index 320ff20b2ef5d..29d45fb42ffcc 100644 --- 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 visitedQueries = new ArrayList<>(); + scriptQueryBuilder.visit(createTestVisitor(visitedQueries)); + assertEquals(2, visitedQueries.size()); + } }