Use QueryVisitor when extracting PercolatorQuery list for highlighting (

#53728)

The highlighting phase for percolator queries currently uses some custom query
traversal logic to find all instances of PercolatorQuery in the query tree for the
current search context. This commit converts things to instead use a QueryVisitor,
which future-proofs us against new wrapper queries or queries from custom
plugins that the percolator module doesn't know about.

modules/percolator/src/main/java/org/elasticsearch/percolator/PercolatorHighlightSubFetchPhase.java

@@ -21,16 +21,11 @@
 
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.ReaderUtil;
-import org.apache.lucene.search.BooleanClause;
-import org.apache.lucene.search.BooleanQuery;
-import org.apache.lucene.search.BoostQuery;
-import org.apache.lucene.search.ConstantScoreQuery;
-import org.apache.lucene.search.DisjunctionMaxQuery;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
+import org.apache.lucene.search.QueryVisitor;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.document.DocumentField;
-import org.elasticsearch.common.lucene.search.function.FunctionScoreQuery;
 import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.fetch.FetchSubPhase;
@@ -139,33 +134,18 @@ public void hitsExecute(SearchContext context, SearchHit[] hits) throws IOExcept
     }
 
     static List<PercolateQuery> locatePercolatorQuery(Query query) {
-        if (query instanceof PercolateQuery) {
-            return Collections.singletonList((PercolateQuery) query);
-        } else if (query instanceof BooleanQuery) {
-            List<PercolateQuery> percolateQueries = new ArrayList<>();
-            for (BooleanClause clause : ((BooleanQuery) query).clauses()) {
-                List<PercolateQuery> result = locatePercolatorQuery(clause.getQuery());
-                if (result.isEmpty() == false) {
-                    percolateQueries.addAll(result);
-                }
-            }
-            return percolateQueries;
-        } else if (query instanceof DisjunctionMaxQuery) {
-            List<PercolateQuery> percolateQueries = new ArrayList<>();
-            for (Query disjunct : ((DisjunctionMaxQuery) query).getDisjuncts()) {
-                List<PercolateQuery> result = locatePercolatorQuery(disjunct);
-                if (result.isEmpty() == false) {
-                    percolateQueries.addAll(result);
+        if (query == null) {
+            return Collections.emptyList();
+        }
+        List<PercolateQuery> queries = new ArrayList<>();
+        query.visit(new QueryVisitor() {
+            @Override
+            public void visitLeaf(Query query) {
+                if (query instanceof PercolateQuery) {
+                    queries.add((PercolateQuery)query);
                 }
             }
-            return  percolateQueries;
-        } else if (query instanceof ConstantScoreQuery) {
-            return locatePercolatorQuery(((ConstantScoreQuery) query).getQuery());
-        } else if (query instanceof BoostQuery) {
-            return locatePercolatorQuery(((BoostQuery) query).getQuery());
-        } else if (query instanceof FunctionScoreQuery) {
-            return locatePercolatorQuery(((FunctionScoreQuery) query).getSubQuery());
-        }
-        return Collections.emptyList();
+        });
+        return queries;
     }
 }

modules/percolator/src/test/java/org/elasticsearch/percolator/PercolatorHighlightSubFetchPhaseTests.java

@@ -37,6 +37,7 @@
 import java.util.Collections;
 
 import static java.util.Collections.emptyMap;
+import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.sameInstance;
@@ -99,8 +100,11 @@ public void testLocatePercolatorQuery() {
         bq.add(percolateQuery, BooleanClause.Occur.FILTER);
         bq.add(percolateQuery2, BooleanClause.Occur.FILTER);
         assertThat(PercolatorHighlightSubFetchPhase.locatePercolatorQuery(bq.build()).size(), equalTo(2));
-        assertThat(PercolatorHighlightSubFetchPhase.locatePercolatorQuery(bq.build()).get(0), sameInstance(percolateQuery));
-        assertThat(PercolatorHighlightSubFetchPhase.locatePercolatorQuery(bq.build()).get(1), sameInstance(percolateQuery2));
+        assertThat(PercolatorHighlightSubFetchPhase.locatePercolatorQuery(bq.build()),
+            containsInAnyOrder(sameInstance(percolateQuery), sameInstance(percolateQuery2)));
+
+        assertNotNull(PercolatorHighlightSubFetchPhase.locatePercolatorQuery(null));
+        assertThat(PercolatorHighlightSubFetchPhase.locatePercolatorQuery(null).size(), equalTo(0));
     }
 
 }