From ddaf3534d908630900f4814806aea1de8f1ebecc Mon Sep 17 00:00:00 2001
From: Alan Woodward <romseygeek@apache.org>
Date: Wed, 18 Mar 2020 15:18:49 +0000
Subject: [PATCH] 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.
---
 .../PercolatorHighlightSubFetchPhase.java     | 44 +++++--------------
 ...PercolatorHighlightSubFetchPhaseTests.java |  8 +++-
 2 files changed, 18 insertions(+), 34 deletions(-)

diff --git a/modules/percolator/src/main/java/org/elasticsearch/percolator/PercolatorHighlightSubFetchPhase.java b/modules/percolator/src/main/java/org/elasticsearch/percolator/PercolatorHighlightSubFetchPhase.java
index d13a54cba0467..cf76f531b6d66 100644
--- a/modules/percolator/src/main/java/org/elasticsearch/percolator/PercolatorHighlightSubFetchPhase.java
+++ b/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;
     }
 }
diff --git a/modules/percolator/src/test/java/org/elasticsearch/percolator/PercolatorHighlightSubFetchPhaseTests.java b/modules/percolator/src/test/java/org/elasticsearch/percolator/PercolatorHighlightSubFetchPhaseTests.java
index 291a42c14665f..09dd2e67e9901 100644
--- a/modules/percolator/src/test/java/org/elasticsearch/percolator/PercolatorHighlightSubFetchPhaseTests.java
+++ b/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));
     }
 
 }