originalQuery is now final

Author: piergm

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

@@ -87,7 +87,7 @@ public void process(HitContext hit) throws IOException {
                             continue;
                         }
 
-                        SearchHighlightContext highlight = new SearchHighlightContext(fetchContext.highlight().fields());
+                        SearchHighlightContext highlight = new SearchHighlightContext(fetchContext.highlight().fields(), null);
                         FetchSubPhaseProcessor processor = highlightPhase.getProcessor(fetchContext, highlight, query);
                         StoredFieldLoader storedFields = StoredFieldLoader.fromSpec(processor.storedFieldsSpec());
                         LeafStoredFieldLoader leafStoredFields = storedFields.getLoader(percolatorLeafReaderContext, null);

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

@@ -45,7 +45,7 @@ public void testHitsExecutionNeeded() {
         );
         PercolatorHighlightSubFetchPhase subFetchPhase = new PercolatorHighlightSubFetchPhase(emptyMap());
         FetchContext fetchContext = mock(FetchContext.class);
-        Mockito.when(fetchContext.highlight()).thenReturn(new SearchHighlightContext(Collections.emptyList()));
+        Mockito.when(fetchContext.highlight()).thenReturn(new SearchHighlightContext(Collections.emptyList(), null));
         Mockito.when(fetchContext.query()).thenReturn(new MatchAllDocsQuery());
 
         assertNull(subFetchPhase.getProcessor(fetchContext));

server/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java

@@ -118,7 +118,7 @@ protected void setupInnerHitsContext(
             }
         }
         if (innerHitBuilder.getHighlightBuilder() != null) {
-            innerHitsContext.highlight(innerHitBuilder.getHighlightBuilder().build(searchExecutionContext));
+            innerHitsContext.highlight(innerHitBuilder.getHighlightBuilder().build(searchExecutionContext, query));
         }
         ParsedQuery parsedQuery = new ParsedQuery(query.toQuery(searchExecutionContext), searchExecutionContext.copyNamedQueries());
         innerHitsContext.parsedQuery(parsedQuery);

server/src/main/java/org/elasticsearch/search/SearchService.java

@@ -90,7 +90,6 @@
 import org.elasticsearch.search.fetch.subphase.FetchFieldsContext;
 import org.elasticsearch.search.fetch.subphase.ScriptFieldsContext.ScriptField;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
-import org.elasticsearch.search.fetch.subphase.highlight.SearchHighlightContext;
 import org.elasticsearch.search.internal.AliasFilter;
 import org.elasticsearch.search.internal.InternalScrollSearchRequest;
 import org.elasticsearch.search.internal.LegacyReaderContext;
@@ -1285,9 +1284,7 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
         if (source.highlighter() != null) {
             HighlightBuilder highlightBuilder = source.highlighter();
             try {
-                SearchHighlightContext searchHighlightContext = highlightBuilder.build(searchExecutionContext);
-                searchHighlightContext.setOriginalQuery(source.query());
-                context.highlight(searchHighlightContext);
+                context.highlight(highlightBuilder.build(searchExecutionContext, source.query()));
             } catch (IOException e) {
                 throw new SearchException(shardTarget, "failed to create SearchContextHighlighter", e);
             }

server/src/main/java/org/elasticsearch/search/aggregations/metrics/TopHitsAggregatorFactory.java

@@ -118,7 +118,9 @@ public Aggregator createInternal(Aggregator parent, CardinalityUpperBound cardin
             subSearchContext.fetchSourceContext(fetchSourceContext);
         }
         if (highlightBuilder != null) {
-            subSearchContext.highlight(highlightBuilder.build(subSearchContext.getSearchExecutionContext()));
+            subSearchContext.highlight(
+                highlightBuilder.build(subSearchContext.getSearchExecutionContext(), highlightBuilder.highlightQuery())
+            );
         }
         return new TopHitsAggregator(subSearchContext, name, context, parent, metadata);
     }

server/src/main/java/org/elasticsearch/search/fetch/subphase/highlight/HighlightBuilder.java

@@ -279,7 +279,7 @@ public static HighlightBuilder fromXContent(XContentParser p) {
         return PARSER.apply(p, new HighlightBuilder());
     }
 
-    public SearchHighlightContext build(SearchExecutionContext context) throws IOException {
+    public SearchHighlightContext build(SearchExecutionContext context, QueryBuilder originalQuery) throws IOException {
         // create template global options that are later merged with any partial field options
         final SearchHighlightContext.FieldOptions.Builder globalOptionsBuilder = new SearchHighlightContext.FieldOptions.Builder();
         globalOptionsBuilder.encoder(this.encoder);
@@ -303,7 +303,7 @@ public SearchHighlightContext build(SearchExecutionContext context) throws IOExc
                 new SearchHighlightContext.Field(field.name(), fieldOptionsBuilder.merge(globalOptionsBuilder.build()).build())
             );
         }
-        return new SearchHighlightContext(fieldOptions);
+        return new SearchHighlightContext(fieldOptions, originalQuery);
     }
 
     /**

server/src/main/java/org/elasticsearch/search/fetch/subphase/highlight/SearchHighlightContext.java

@@ -22,14 +22,15 @@
 
 public class SearchHighlightContext {
     private final Map<String, Field> fields;
-    private QueryBuilder originalQuery;
+    private final QueryBuilder originalQuery;
 
-    public SearchHighlightContext(Collection<Field> fields) {
+    public SearchHighlightContext(Collection<Field> fields, QueryBuilder originalQuery) {
         assert fields != null;
         this.fields = Maps.newLinkedHashMapWithExpectedSize(fields.size());
         for (Field field : fields) {
             this.fields.put(field.field, field);
         }
+        this.originalQuery = originalQuery;
     }
 
     public Collection<Field> fields() {
@@ -40,10 +41,6 @@ public QueryBuilder originalQuery() {
         return originalQuery;
     }
 
-    public void setOriginalQuery(QueryBuilder originalQuery) {
-        this.originalQuery = originalQuery;
-    }
-
     public static class Field {
         private final String field;
         private final FieldOptions fieldOptions;

server/src/test/java/org/elasticsearch/search/fetch/subphase/highlight/HighlightBuilderTests.java

@@ -328,7 +328,7 @@ public MappedFieldType getFieldType(String name) {
         for (int runs = 0; runs < NUMBER_OF_TESTBUILDERS; runs++) {
             HighlightBuilder highlightBuilder = randomHighlighterBuilder();
             highlightBuilder = Rewriteable.rewrite(highlightBuilder, mockContext);
-            SearchHighlightContext highlight = highlightBuilder.build(mockContext);
+            SearchHighlightContext highlight = highlightBuilder.build(mockContext, highlightBuilder.highlightQuery);
             for (SearchHighlightContext.Field field : highlight.fields()) {
                 String encoder = highlightBuilder.encoder() != null ? highlightBuilder.encoder() : HighlightBuilder.DEFAULT_ENCODER;
                 assertEquals(encoder, field.fieldOptions().encoder());

test/framework/src/main/java/org/elasticsearch/search/fetch/HighlighterTestCase.java

@@ -109,7 +109,7 @@ protected static void assertHighlights(Map<String, HighlightField> highlights, S
 
     private static FetchContext fetchContext(SearchExecutionContext context, SearchSourceBuilder search) throws IOException {
         FetchContext fetchContext = mock(FetchContext.class);
-        when(fetchContext.highlight()).thenReturn(search.highlighter().build(context));
+        when(fetchContext.highlight()).thenReturn(search.highlighter().build(context, search.query()));
         when(fetchContext.parsedQuery()).thenReturn(new ParsedQuery(search.query().toQuery(context)));
         when(fetchContext.getSearchExecutionContext()).thenReturn(context);
         when(fetchContext.sourceLoader()).thenReturn(context.newSourceLoader(false));