From 4b736d2e0c01d55eb320c205fc9c90da52dab910 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Christoph=20B=C3=BCscher?= <christoph@elastic.co>
Date: Mon, 8 Feb 2016 14:26:16 +0100
Subject: [PATCH] Changing parse signature in QueryParseElement

---
 .../action/suggest/TransportSuggestAction.java     |  5 +++--
 .../search/suggest/SuggestContextParser.java       |  5 ++---
 .../search/suggest/SuggestParseElement.java        | 14 ++++++++------
 .../search/suggest/SuggestionSearchContext.java    | 10 ++++++++++
 .../completion/CompletionSuggestParser.java        |  8 ++++----
 .../suggest/completion/CompletionSuggester.java    |  4 +++-
 .../completion/CompletionSuggestionContext.java    |  9 +--------
 .../search/suggest/phrase/PhraseSuggestParser.java |  5 +++--
 .../search/suggest/phrase/PhraseSuggester.java     |  3 ++-
 .../suggest/phrase/PhraseSuggestionContext.java    | 12 +-----------
 .../search/suggest/term/TermSuggestParser.java     |  5 +++--
 .../suggest/AbstractSuggestionBuilderTestCase.java |  2 +-
 .../search/suggest/CustomSuggester.java            |  2 +-
 13 files changed, 42 insertions(+), 42 deletions(-)

diff --git a/core/src/main/java/org/elasticsearch/action/suggest/TransportSuggestAction.java b/core/src/main/java/org/elasticsearch/action/suggest/TransportSuggestAction.java
index 0ed98578557..36fb079d257 100644
--- a/core/src/main/java/org/elasticsearch/action/suggest/TransportSuggestAction.java
+++ b/core/src/main/java/org/elasticsearch/action/suggest/TransportSuggestAction.java
@@ -39,6 +39,7 @@ import org.elasticsearch.common.xcontent.XContentFactory;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.index.IndexService;
 import org.elasticsearch.index.engine.Engine;
+import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.index.shard.IndexShard;
 import org.elasticsearch.index.suggest.stats.ShardSuggestMetric;
 import org.elasticsearch.indices.IndicesService;
@@ -131,6 +132,7 @@ public class TransportSuggestAction extends TransportBroadcastAction<SuggestRequ
     protected ShardSuggestResponse shardOperation(ShardSuggestRequest request) {
         IndexService indexService = indicesService.indexServiceSafe(request.shardId().getIndex());
         IndexShard indexShard = indexService.getShard(request.shardId().id());
+        QueryShardContext shardContext = indexShard.getQueryShardContext();
         ShardSuggestMetric suggestMetric = indexShard.getSuggestMetric();
         suggestMetric.preSuggest();
         long startTime = System.nanoTime();
@@ -142,8 +144,7 @@ public class TransportSuggestAction extends TransportBroadcastAction<SuggestRequ
                 if (parser.nextToken() != XContentParser.Token.START_OBJECT) {
                     throw new IllegalArgumentException("suggest content missing");
                 }
-                final SuggestionSearchContext context = suggestPhase.parseElement().parseInternal(parser, indexService.mapperService(),
-                        indexService.fieldData(), request.shardId().getIndexName(), request.shardId().id());
+                final SuggestionSearchContext context = suggestPhase.parseElement().parseInternal(parser, indexShard.getQueryShardContext());
                 final Suggest result = suggestPhase.execute(context, searcher.searcher());
                 return new ShardSuggestResponse(request.shardId(), result);
             }
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/SuggestContextParser.java b/core/src/main/java/org/elasticsearch/search/suggest/SuggestContextParser.java
index a7aa3fd60b6..53d510bf530 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/SuggestContextParser.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/SuggestContextParser.java
@@ -19,12 +19,11 @@
 package org.elasticsearch.search.suggest;
 
 import org.elasticsearch.common.xcontent.XContentParser;
-import org.elasticsearch.index.fielddata.IndexFieldDataService;
-import org.elasticsearch.index.mapper.MapperService;
+import org.elasticsearch.index.query.QueryShardContext;
 
 import java.io.IOException;
 
 public interface SuggestContextParser {
-    SuggestionSearchContext.SuggestionContext parse(XContentParser parser, MapperService mapperService, IndexFieldDataService indexFieldDataService) throws IOException;
+    SuggestionSearchContext.SuggestionContext parse(XContentParser parser, QueryShardContext shardContext) throws IOException;
 
 }
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/SuggestParseElement.java b/core/src/main/java/org/elasticsearch/search/suggest/SuggestParseElement.java
index 52b728a476f..216406e9e0b 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/SuggestParseElement.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/SuggestParseElement.java
@@ -21,8 +21,8 @@ package org.elasticsearch.search.suggest;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.xcontent.XContentParser;
-import org.elasticsearch.index.fielddata.IndexFieldDataService;
 import org.elasticsearch.index.mapper.MapperService;
+import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.search.SearchParseElement;
 import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.suggest.SuggestionSearchContext.SuggestionContext;
@@ -44,14 +44,13 @@ public final class SuggestParseElement implements SearchParseElement {
 
     @Override
     public void parse(XContentParser parser, SearchContext context) throws Exception {
-        SuggestionSearchContext suggestionSearchContext = parseInternal(parser, context.mapperService(), context.fieldData(),
-                context.shardTarget().index(), context.shardTarget().shardId());
+        SuggestionSearchContext suggestionSearchContext = parseInternal(parser, context.indexShard().getQueryShardContext());
         context.suggest(suggestionSearchContext);
     }
 
-    public SuggestionSearchContext parseInternal(XContentParser parser, MapperService mapperService, IndexFieldDataService fieldDataService,
-                                                 String index, int shardId) throws IOException {
+    public SuggestionSearchContext parseInternal(XContentParser parser, QueryShardContext shardContext) throws IOException {
         SuggestionSearchContext suggestionSearchContext = new SuggestionSearchContext();
+        MapperService mapperService = shardContext.getMapperService();
 
         BytesRef globalText = null;
         String fieldName = null;
@@ -95,7 +94,7 @@ public final class SuggestParseElement implements SearchParseElement {
                             throw new IllegalArgumentException("Suggester[" + fieldName + "] not supported");
                         }
                         final SuggestContextParser contextParser = suggesters.get(fieldName).getContextParser();
-                        suggestionContext = contextParser.parse(parser, mapperService, fieldDataService);
+                        suggestionContext = contextParser.parse(parser, shardContext);
                     }
                 }
                 if (suggestionContext != null) {
@@ -109,7 +108,10 @@ public final class SuggestParseElement implements SearchParseElement {
                         suggestionContext.setRegex(regex);
                         suggestionContext.setText(regex);
                     }
+                    suggestionContext.setShardContext(shardContext);
                     suggestionContexts.put(suggestionName, suggestionContext);
+                } else {
+                    throw new IllegalArgumentException("suggestion context could not be parsed correctly");
                 }
             }
         }
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/SuggestionSearchContext.java b/core/src/main/java/org/elasticsearch/search/suggest/SuggestionSearchContext.java
index b662df33f28..02953356876 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/SuggestionSearchContext.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/SuggestionSearchContext.java
@@ -20,6 +20,7 @@ package org.elasticsearch.search.suggest;
 
 import org.apache.lucene.analysis.Analyzer;
 import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.index.query.QueryShardContext;
 
 import java.util.LinkedHashMap;
 import java.util.Map;
@@ -47,6 +48,7 @@ public class SuggestionSearchContext {
         private Analyzer analyzer;
         private int size = 5;
         private int shardSize = -1;
+        private QueryShardContext shardContext;
 
         public BytesRef getText() {
             return text;
@@ -117,6 +119,14 @@ public class SuggestionSearchContext {
             }
             this.shardSize = shardSize;
         }
+
+        public void setShardContext(QueryShardContext context) {
+            this.shardContext = context;
+        }
+
+        public QueryShardContext getShardContext() {
+            return this.shardContext;
+        }
     }
 
 }
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestParser.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestParser.java
index 702b03f359e..982c86957b6 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestParser.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestParser.java
@@ -20,17 +20,16 @@ package org.elasticsearch.search.suggest.completion;
 
 import org.apache.lucene.analysis.Analyzer;
 import org.elasticsearch.ElasticsearchException;
-import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.unit.Fuzziness;
 import org.elasticsearch.common.xcontent.ObjectParser;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentFactory;
 import org.elasticsearch.common.xcontent.XContentParser;
-import org.elasticsearch.index.fielddata.IndexFieldDataService;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.core.CompletionFieldMapper;
+import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.index.query.RegexpFlag;
 import org.elasticsearch.search.suggest.SuggestContextParser;
 import org.elasticsearch.search.suggest.SuggestUtils.Fields;
@@ -135,8 +134,9 @@ public class CompletionSuggestParser implements SuggestContextParser {
     }
 
     @Override
-    public SuggestionSearchContext.SuggestionContext parse(XContentParser parser, MapperService mapperService, IndexFieldDataService fieldDataService) throws IOException {
-        final CompletionSuggestionContext suggestion = new CompletionSuggestionContext(completionSuggester, mapperService, fieldDataService);
+    public SuggestionSearchContext.SuggestionContext parse(XContentParser parser,  QueryShardContext shardContext) throws IOException {
+        MapperService mapperService = shardContext.getMapperService();
+        final CompletionSuggestionContext suggestion = new CompletionSuggestionContext(completionSuggester, mapperService);
         final ContextAndSuggest contextAndSuggest = new ContextAndSuggest(mapperService);
         TLP_PARSER.parse(parser, suggestion, contextAndSuggest);
         final XContentParser contextParser = contextAndSuggest.contextParser;
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggester.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggester.java
index 86f3b87b059..5b8367ef97c 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggester.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggester.java
@@ -35,6 +35,7 @@ import org.elasticsearch.index.fielddata.AtomicFieldData;
 import org.elasticsearch.index.fielddata.ScriptDocValues;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.core.CompletionFieldMapper;
+import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.search.suggest.Suggest;
 import org.elasticsearch.search.suggest.SuggestContextParser;
 import org.elasticsearch.search.suggest.Suggester;
@@ -90,7 +91,8 @@ public class CompletionSuggester extends Suggester<CompletionSuggestionContext>
                 for (String field : payloadFields) {
                     MappedFieldType payloadFieldType = suggestionContext.getMapperService().fullName(field);
                     if (payloadFieldType != null) {
-                        final AtomicFieldData data = suggestionContext.getIndexFieldDataService().getForField(payloadFieldType)
+                        QueryShardContext shardContext = suggestionContext.getShardContext();
+                        final AtomicFieldData data = shardContext.getForField(payloadFieldType)
                                 .load(subReaderContext);
                         final ScriptDocValues scriptValues = data.getScriptValues();
                         scriptValues.setNextDocId(subDocId);
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionContext.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionContext.java
index 535151b476e..328005a8cb2 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionContext.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionContext.java
@@ -20,7 +20,6 @@ package org.elasticsearch.search.suggest.completion;
 
 import org.apache.lucene.search.suggest.document.CompletionQuery;
 import org.elasticsearch.common.unit.Fuzziness;
-import org.elasticsearch.index.fielddata.IndexFieldDataService;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.core.CompletionFieldMapper;
 import org.elasticsearch.search.suggest.Suggester;
@@ -44,12 +43,10 @@ public class CompletionSuggestionContext extends SuggestionSearchContext.Suggest
     private CompletionSuggestionBuilder.RegexOptionsBuilder regexOptionsBuilder;
     private Map<String, List<ContextMapping.QueryContext>> queryContexts = Collections.emptyMap();
     private final MapperService mapperService;
-    private final IndexFieldDataService indexFieldDataService;
     private Set<String> payloadFields = Collections.emptySet();
 
-    CompletionSuggestionContext(Suggester suggester, MapperService mapperService, IndexFieldDataService indexFieldDataService) {
+    CompletionSuggestionContext(Suggester suggester, MapperService mapperService) {
         super(suggester);
-        this.indexFieldDataService = indexFieldDataService;
         this.mapperService = mapperService;
     }
 
@@ -78,10 +75,6 @@ public class CompletionSuggestionContext extends SuggestionSearchContext.Suggest
         return mapperService;
     }
 
-    IndexFieldDataService getIndexFieldDataService() {
-        return indexFieldDataService;
-    }
-
     void setPayloadFields(Set<String> fields) {
         this.payloadFields = fields;
     }
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggestParser.java b/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggestParser.java
index eb7254c7226..b6739866fa7 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggestParser.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggestParser.java
@@ -26,9 +26,9 @@ import org.elasticsearch.common.ParseFieldMatcher;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.common.xcontent.XContentParser.Token;
 import org.elasticsearch.index.analysis.ShingleTokenFilterFactory;
-import org.elasticsearch.index.fielddata.IndexFieldDataService;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.MapperService;
+import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.script.CompiledScript;
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.Template;
@@ -51,7 +51,8 @@ public final class PhraseSuggestParser implements SuggestContextParser {
     }
 
     @Override
-    public SuggestionSearchContext.SuggestionContext parse(XContentParser parser, MapperService mapperService, IndexFieldDataService fieldDataService) throws IOException {
+    public SuggestionSearchContext.SuggestionContext parse(XContentParser parser, QueryShardContext shardContext) throws IOException {
+        MapperService mapperService = shardContext.getMapperService();
         PhraseSuggestionContext suggestion = new PhraseSuggestionContext(suggester);
         ParseFieldMatcher parseFieldMatcher = mapperService.getIndexSettings().getParseFieldMatcher();
         XContentParser.Token token;
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggester.java b/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggester.java
index 74dcb33e643..d4731568b59 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggester.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggester.java
@@ -42,6 +42,7 @@ import org.elasticsearch.search.suggest.SuggestContextParser;
 import org.elasticsearch.search.suggest.SuggestUtils;
 import org.elasticsearch.search.suggest.Suggester;
 import org.elasticsearch.search.suggest.SuggestionBuilder;
+import org.elasticsearch.search.suggest.SuggestionSearchContext.SuggestionContext;
 import org.elasticsearch.search.suggest.phrase.NoisyChannelSpellChecker.Result;
 
 import java.io.IOException;
@@ -143,7 +144,7 @@ public final class PhraseSuggester extends Suggester<PhraseSuggestionContext> {
         return response;
     }
 
-    private PhraseSuggestion.Entry buildResultEntry(PhraseSuggestionContext suggestion, CharsRefBuilder spare, double cutoffScore) {
+    private PhraseSuggestion.Entry buildResultEntry(SuggestionContext suggestion, CharsRefBuilder spare, double cutoffScore) {
         spare.copyUTF8Bytes(suggestion.getText());
         return new PhraseSuggestion.Entry(new Text(spare.toString()), 0, spare.length(), cutoffScore);
     }
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggestionContext.java b/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggestionContext.java
index 4b63c03f8ec..cf5eb4febab 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggestionContext.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggestionContext.java
@@ -20,7 +20,6 @@ package org.elasticsearch.search.suggest.phrase;
 
 import org.apache.lucene.analysis.Analyzer;
 import org.apache.lucene.util.BytesRef;
-import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.script.CompiledScript;
 import org.elasticsearch.search.suggest.DirectSpellcheckerSettings;
 import org.elasticsearch.search.suggest.Suggester;
@@ -31,7 +30,7 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
-class PhraseSuggestionContext extends SuggestionContext {
+public class PhraseSuggestionContext extends SuggestionContext {
     static final boolean DEFAULT_COLLATE_PRUNE = false;
     static final boolean DEFAULT_REQUIRE_UNIGRAM = true;
     static final float DEFAULT_CONFIDENCE = 1.0f;
@@ -54,7 +53,6 @@ class PhraseSuggestionContext extends SuggestionContext {
     private List<DirectCandidateGenerator> generators = new ArrayList<>();
     private Map<String, Object> collateScriptParams = new HashMap<>(1);
     private WordScorer.WordScorerFactory scorer;
-    private QueryShardContext shardContext;
 
     public PhraseSuggestionContext(Suggester<? extends PhraseSuggestionContext> suggester) {
         super(suggester);
@@ -216,12 +214,4 @@ class PhraseSuggestionContext extends SuggestionContext {
         return prune;
     }
 
-    public void setShardContext(QueryShardContext context) {
-        this.shardContext = context;
-    }
-
-    public QueryShardContext getShardContext() {
-        return this.shardContext;
-    }
-
 }
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/term/TermSuggestParser.java b/core/src/main/java/org/elasticsearch/search/suggest/term/TermSuggestParser.java
index a2fd680c215..f819ed412af 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/term/TermSuggestParser.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/term/TermSuggestParser.java
@@ -20,8 +20,8 @@ package org.elasticsearch.search.suggest.term;
 
 import org.elasticsearch.common.ParseFieldMatcher;
 import org.elasticsearch.common.xcontent.XContentParser;
-import org.elasticsearch.index.fielddata.IndexFieldDataService;
 import org.elasticsearch.index.mapper.MapperService;
+import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.search.suggest.DirectSpellcheckerSettings;
 import org.elasticsearch.search.suggest.SuggestContextParser;
 import org.elasticsearch.search.suggest.SuggestUtils;
@@ -38,7 +38,8 @@ public final class TermSuggestParser implements SuggestContextParser {
     }
 
     @Override
-    public SuggestionSearchContext.SuggestionContext parse(XContentParser parser, MapperService mapperService, IndexFieldDataService fieldDataService) throws IOException {
+    public SuggestionSearchContext.SuggestionContext parse(XContentParser parser, QueryShardContext shardContext) throws IOException {
+        MapperService mapperService = shardContext.getMapperService();
         XContentParser.Token token;
         String fieldName = null;
         TermSuggestionContext suggestion = new TermSuggestionContext(suggester);
diff --git a/core/src/test/java/org/elasticsearch/search/suggest/AbstractSuggestionBuilderTestCase.java b/core/src/test/java/org/elasticsearch/search/suggest/AbstractSuggestionBuilderTestCase.java
index 501041c6af2..0df3abecefd 100644
--- a/core/src/test/java/org/elasticsearch/search/suggest/AbstractSuggestionBuilderTestCase.java
+++ b/core/src/test/java/org/elasticsearch/search/suggest/AbstractSuggestionBuilderTestCase.java
@@ -271,7 +271,7 @@ public abstract class AbstractSuggestionBuilderTestCase<SB extends SuggestionBui
 
             XContentParser parser = XContentHelper.createParser(xContentBuilder.bytes());
             parser.nextToken(); // set cursor to START_OBJECT
-            SuggestionSearchContext suggestionSearchContext = parseElement.parseInternal(parser, mockMapperService, null, "test", 1);
+            SuggestionSearchContext suggestionSearchContext = parseElement.parseInternal(parser, mockShardContext);
             SuggestionContext oldSchoolContext = suggestionSearchContext.suggestions().get(suggestionBuilder.name());
 
             SuggestionContext newSchoolContext = suggestionBuilder.build(mockShardContext);
diff --git a/core/src/test/java/org/elasticsearch/search/suggest/CustomSuggester.java b/core/src/test/java/org/elasticsearch/search/suggest/CustomSuggester.java
index 4dbae08080a..e2ca13003bb 100644
--- a/core/src/test/java/org/elasticsearch/search/suggest/CustomSuggester.java
+++ b/core/src/test/java/org/elasticsearch/search/suggest/CustomSuggester.java
@@ -54,7 +54,7 @@ public class CustomSuggester extends Suggester<CustomSuggester.CustomSuggestions
 
     @Override
     public SuggestContextParser getContextParser() {
-        return (parser, mapperService, fieldData) -> {
+        return (parser, shardContext) -> {
             Map<String, Object> options = parser.map();
             CustomSuggestionsContext suggestionContext = new CustomSuggestionsContext(CustomSuggester.this, options);
             suggestionContext.setField((String) options.get("field"));