Make fetch sub phases pluggable

core/src/main/java/org/elasticsearch/percolator/PercolateContext.java

@@ -32,6 +32,7 @@
 import org.elasticsearch.action.percolate.PercolateShardRequest;
 import org.elasticsearch.action.search.SearchType;
 import org.elasticsearch.cache.recycler.PageCacheRecycler;
+import org.elasticsearch.client.Client;
 import org.elasticsearch.common.*;
 import org.elasticsearch.common.collect.ImmutableOpenMap;
 import org.elasticsearch.common.lease.Releasables;
@@ -58,7 +59,7 @@
 import org.elasticsearch.search.dfs.DfsSearchResult;
 import org.elasticsearch.search.fetch.FetchSearchResult;
 import org.elasticsearch.search.fetch.FetchSubPhase;
-import org.elasticsearch.search.fetch.fielddata.FieldDataFieldsContext;
+import org.elasticsearch.search.fetch.FetchSubPhaseContext;
 import org.elasticsearch.search.fetch.innerhits.InnerHitsContext;
 import org.elasticsearch.search.fetch.script.ScriptFieldsContext;
 import org.elasticsearch.search.fetch.source.FetchSourceContext;
@@ -116,6 +117,7 @@ public class PercolateContext extends SearchContext {
     private SearchContextAggregations aggregations;
     private QuerySearchResult querySearchResult;
     private Sort sort;
+    private final Map<String, FetchSubPhaseContext> subPhaseContexts = new HashMap<>();
 
     public PercolateContext(PercolateShardRequest request, SearchShardTarget searchShardTarget, IndexShard indexShard,
                             IndexService indexService, PageCacheRecycler pageCacheRecycler,
@@ -282,6 +284,15 @@ public SearchContext aggregations(SearchContextAggregations aggregations) {
         return this;
     }
 
+    @Override
+    public <SubPhaseContext extends FetchSubPhaseContext> SubPhaseContext getFetchSubPhaseContext(FetchSubPhase.ContextFactory<SubPhaseContext> contextFactory) {
+        String subPhaseName = contextFactory.getName();
+        if (subPhaseContexts.get(subPhaseName) == null) {
+            subPhaseContexts.put(subPhaseName, contextFactory.newContextInstance());
+        }
+        return (SubPhaseContext) subPhaseContexts.get(subPhaseName);
+    }
+
     // Unused:
     @Override
     public void preProcess() {
@@ -378,16 +389,6 @@ public void addRescore(RescoreSearchContext rescore) {
         throw new UnsupportedOperationException();
     }
 
-    @Override
-    public boolean hasFieldDataFields() {
-        throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public FieldDataFieldsContext fieldDataFields() {
-        throw new UnsupportedOperationException();
-    }
-
     @Override
     public boolean hasScriptFields() {
         throw new UnsupportedOperationException();

core/src/main/java/org/elasticsearch/search/SearchModule.java

@@ -32,6 +32,7 @@
 import org.elasticsearch.search.controller.SearchPhaseController;
 import org.elasticsearch.search.dfs.DfsPhase;
 import org.elasticsearch.search.fetch.FetchPhase;
+import org.elasticsearch.search.fetch.FetchSubPhaseModule;
 import org.elasticsearch.search.fetch.explain.ExplainFetchSubPhase;
 import org.elasticsearch.search.fetch.fielddata.FieldDataFieldsFetchSubPhase;
 import org.elasticsearch.search.fetch.innerhits.InnerHitsFetchSubPhase;
@@ -63,7 +64,8 @@ public Iterable<? extends Module> spawnModules() {
                 new HighlightModule(),
                 new SuggestModule(),
                 new FunctionScoreModule(),
-                new AggregationModule());
+                new AggregationModule(),
+                new FetchSubPhaseModule());
     }
 
     @Override
@@ -73,14 +75,6 @@ protected void configure() {
         bind(SearchPhaseController.class).asEagerSingleton();
 
         bind(FetchPhase.class).asEagerSingleton();
-        bind(ExplainFetchSubPhase.class).asEagerSingleton();
-        bind(FieldDataFieldsFetchSubPhase.class).asEagerSingleton();
-        bind(ScriptFieldsFetchSubPhase.class).asEagerSingleton();
-        bind(FetchSourceSubPhase.class).asEagerSingleton();
-        bind(VersionFetchSubPhase.class).asEagerSingleton();
-        bind(MatchedQueriesFetchSubPhase.class).asEagerSingleton();
-        bind(HighlightPhase.class).asEagerSingleton();
-        bind(InnerHitsFetchSubPhase.class).asEagerSingleton();
 
         bind(SearchServiceTransportAction.class).asEagerSingleton();
         bind(MoreLikeThisFetchService.class).asEagerSingleton();

core/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java

@@ -45,6 +45,7 @@
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.internal.SourceFieldMapper;
 import org.elasticsearch.index.mapper.object.ObjectMapper;
+import org.elasticsearch.index.query.functionscore.ScoreFunctionParser;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.SearchHitField;
 import org.elasticsearch.search.SearchParseElement;
@@ -83,13 +84,10 @@ public class FetchPhase implements SearchPhase {
     private final FetchSubPhase[] fetchSubPhases;
 
     @Inject
-    public FetchPhase(HighlightPhase highlightPhase, ScriptFieldsFetchSubPhase scriptFieldsPhase,
-                      MatchedQueriesFetchSubPhase matchedQueriesPhase, ExplainFetchSubPhase explainPhase, VersionFetchSubPhase versionPhase,
-                      FetchSourceSubPhase fetchSourceSubPhase, FieldDataFieldsFetchSubPhase fieldDataFieldsFetchSubPhase, 
-                      InnerHitsFetchSubPhase innerHitsFetchSubPhase) {
+    public FetchPhase(Set<FetchSubPhase> fetchSubPhases, InnerHitsFetchSubPhase innerHitsFetchSubPhase) {
         innerHitsFetchSubPhase.setFetchPhase(this);
-        this.fetchSubPhases = new FetchSubPhase[]{scriptFieldsPhase, matchedQueriesPhase, explainPhase, highlightPhase,
-                fetchSourceSubPhase, versionPhase, fieldDataFieldsFetchSubPhase, innerHitsFetchSubPhase};
+        this.fetchSubPhases = fetchSubPhases.toArray(new FetchSubPhase[fetchSubPhases.size() + 1]);
+        this.fetchSubPhases[fetchSubPhases.size()] = innerHitsFetchSubPhase;
     }
 
     @Override

core/src/main/java/org/elasticsearch/search/fetch/FetchSubPhase.java

@@ -114,4 +114,23 @@ public String getSourcePath(String sourcePath) {
     boolean hitsExecutionNeeded(SearchContext context);
 
     void hitsExecute(SearchContext context, InternalSearchHit[] hits);
+
+    /**
+     * This interface is in the fetch phase plugin mechanism.
+     * Whenever a new search is executed we create a new {@link SearchContext} that holds individual contexts for each {@link org.elasticsearch.search.fetch.FetchSubPhase}.
+     * Fetch phases that use the plugin mechanism must provide a ContextFactory to the SearchContext that creates the fetch phase context and also associates them with a name.
+     * See {@link SearchContext#getFetchSubPhaseContext(FetchSubPhase.ContextFactory)}
+     */
+    public interface ContextFactory<SubPhaseContext extends FetchSubPhaseContext> {
+
+        /**
+         * The name of the context.
+         */
+        public String getName();
+
+        /**
+         * Creates a new instance of a FetchSubPhaseContext that holds all information a FetchSubPhase needs to execute on hits.
+         */
+        public SubPhaseContext newContextInstance();
+    }
 }

core/src/main/java/org/elasticsearch/search/fetch/FetchSubPhaseContext.java

@@ -0,0 +1,47 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.search.fetch;
+
+/**
+ * All configuration and context needed by the FetchSubPhase to execute on hits.
+ * The only required information in this base class is whether or not the sub phase needs to be run at all.
+ * It can be extended by FetchSubPhases to hold information the phase needs to execute on hits.
+ * See {@link org.elasticsearch.search.fetch.FetchSubPhase.ContextFactory} and also {@link org.elasticsearch.search.fetch.fielddata.FieldDataFieldsContext} for an example.
+ */
+public class FetchSubPhaseContext {
+
+    // This is to store if the FetchSubPhase should be executed at all.
+    private boolean hitExecutionNeeded = false;
+
+    /**
+     * Set if this phase should be executed at all.
+     */
+    void setHitExecutionNeeded(boolean hitExecutionNeeded) {
+        this.hitExecutionNeeded = hitExecutionNeeded;
+    }
+
+    /**
+     * Returns if this phase be executed at all.
+     */
+    public boolean hitExecutionNeeded() {
+        return hitExecutionNeeded;
+    }
+
+}

core/src/main/java/org/elasticsearch/search/fetch/FetchSubPhaseModule.java

@@ -0,0 +1,74 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.search.fetch;
+
+import com.google.common.collect.Lists;
+import org.elasticsearch.common.inject.AbstractModule;
+import org.elasticsearch.common.inject.multibindings.Multibinder;
+import org.elasticsearch.search.fetch.explain.ExplainFetchSubPhase;
+import org.elasticsearch.search.fetch.fielddata.FieldDataFieldsFetchSubPhase;
+import org.elasticsearch.search.fetch.innerhits.InnerHitsFetchSubPhase;
+import org.elasticsearch.search.fetch.matchedqueries.MatchedQueriesFetchSubPhase;
+import org.elasticsearch.search.fetch.script.ScriptFieldsFetchSubPhase;
+import org.elasticsearch.search.fetch.source.FetchSourceSubPhase;
+import org.elasticsearch.search.fetch.version.VersionFetchSubPhase;
+import org.elasticsearch.search.highlight.HighlightPhase;
+
+import java.util.List;
+
+/**
+ * Module for registering fetch sub phases. Fetch phases are executed when the document is finally
+ * retrieved from the shard. To implement a new fetch phase one needs to implement the following classes and interfaces
+ * <p/>
+ * <ul>
+ * <li>  {@link FetchSubPhaseParseElement} </li>
+ * <li>  {@link FetchSubPhase} </li>
+ * <li>  {@link FetchSubPhaseContext} </li>
+ * </ul>
+ * <p/>
+ * The FetchSubPhase must then be registered with this module with {@link FetchSubPhaseModule#registerFetchSubPhase(Class<? extends FetchSubPhase>)}.
+ * See {@link FieldDataFieldsFetchSubPhase} for an example.
+ */
+public class FetchSubPhaseModule extends AbstractModule {
+
+    private List<Class<? extends FetchSubPhase>> fetchSubPhases = Lists.newArrayList();
+
+    public FetchSubPhaseModule() {
+        registerFetchSubPhase(ExplainFetchSubPhase.class);
+        registerFetchSubPhase(FieldDataFieldsFetchSubPhase.class);
+        registerFetchSubPhase(ScriptFieldsFetchSubPhase.class);
+        registerFetchSubPhase(FetchSourceSubPhase.class);
+        registerFetchSubPhase(VersionFetchSubPhase.class);
+        registerFetchSubPhase(MatchedQueriesFetchSubPhase.class);
+        registerFetchSubPhase(HighlightPhase.class);
+    }
+
+    public void registerFetchSubPhase(Class<? extends FetchSubPhase> subPhase) {
+        fetchSubPhases.add(subPhase);
+    }
+
+    @Override
+    protected void configure() {
+        Multibinder<FetchSubPhase> parserMapBinder = Multibinder.newSetBinder(binder(), FetchSubPhase.class);
+        for (Class<? extends FetchSubPhase> clazz : fetchSubPhases) {
+            parserMapBinder.addBinding().to(clazz);
+        }
+        bind(InnerHitsFetchSubPhase.class).asEagerSingleton();
+    }
+}

core/src/main/java/org/elasticsearch/search/fetch/FetchSubPhaseParseElement.java

@@ -0,0 +1,48 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.search.fetch;
+
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.search.SearchParseElement;
+import org.elasticsearch.search.internal.SearchContext;
+
+/**
+ * A parse element for a {@link org.elasticsearch.search.fetch.FetchSubPhase} that is used when parsing a search request.
+ */
+public abstract class FetchSubPhaseParseElement<SubPhaseContext extends FetchSubPhaseContext> implements SearchParseElement {
+
+    @Override
+    final public void parse(XContentParser parser, SearchContext context) throws Exception {
+        SubPhaseContext fetchSubPhaseContext = context.getFetchSubPhaseContext(getContextFactory());
+        // this is to make sure that the SubFetchPhase knows it should execute
+        fetchSubPhaseContext.setHitExecutionNeeded(true);
+        innerParse(parser, fetchSubPhaseContext, context);
+    }
+
+    /**
+     * Implement the actual parsing here.
+     */
+    protected abstract void innerParse(XContentParser parser, SubPhaseContext fetchSubPhaseContext, SearchContext searchContext) throws Exception;
+
+    /**
+     * Return the ContextFactory for this FetchSubPhase.
+     */
+    protected abstract FetchSubPhase.ContextFactory<SubPhaseContext> getContextFactory();
+}

core/src/main/java/org/elasticsearch/search/fetch/fielddata/FieldDataFieldsContext.java

@@ -19,13 +19,14 @@
 package org.elasticsearch.search.fetch.fielddata;
 
 import com.google.common.collect.Lists;
+import org.elasticsearch.search.fetch.FetchSubPhaseContext;
 
 import java.util.List;
 
 /**
  * All the required context to pull a field from the field data cache.
  */
-public class FieldDataFieldsContext {
+public class FieldDataFieldsContext extends FetchSubPhaseContext {
 
     public static class FieldDataField {
         private final String name;

core/src/main/java/org/elasticsearch/search/fetch/fielddata/FieldDataFieldsFetchSubPhase.java

@@ -27,6 +27,7 @@
 import org.elasticsearch.search.SearchHitField;
 import org.elasticsearch.search.SearchParseElement;
 import org.elasticsearch.search.fetch.FetchSubPhase;
+import org.elasticsearch.search.fetch.FetchSubPhaseContext;
 import org.elasticsearch.search.internal.InternalSearchHit;
 import org.elasticsearch.search.internal.InternalSearchHitField;
 import org.elasticsearch.search.internal.SearchContext;
@@ -43,6 +44,20 @@
  */
 public class FieldDataFieldsFetchSubPhase implements FetchSubPhase {
 
+    public static final String[] NAMES = {"fielddata_fields", "fielddataFields"};
+    public static final ContextFactory<FieldDataFieldsContext> CONTEXT_FACTORY = new ContextFactory<FieldDataFieldsContext>() {
+
+        @Override
+        public String getName() {
+            return NAMES[0];
+        }
+
+        @Override
+        public FieldDataFieldsContext newContextInstance() {
+            return new FieldDataFieldsContext();
+        }
+    };
+
     @Inject
     public FieldDataFieldsFetchSubPhase() {
     }
@@ -66,12 +81,12 @@ public void hitsExecute(SearchContext context, InternalSearchHit[] hits) {
 
     @Override
     public boolean hitExecutionNeeded(SearchContext context) {
-        return context.hasFieldDataFields();
+        return context.getFetchSubPhaseContext(CONTEXT_FACTORY).hitExecutionNeeded();
     }
 
     @Override
     public void hitExecute(SearchContext context, HitContext hitContext) {
-        for (FieldDataFieldsContext.FieldDataField field : context.fieldDataFields().fields()) {
+        for (FieldDataFieldsContext.FieldDataField field : context.getFetchSubPhaseContext(CONTEXT_FACTORY).fields()) {
             if (hitContext.hit().fieldsOrNull() == null) {
                 hitContext.hit().fields(new HashMap<String, SearchHitField>(2));
             }