Scripting: Move context definitions to instance type classes

core/src/main/java/org/elasticsearch/action/update/UpdateHelper.java

@@ -300,7 +300,7 @@ Result prepareUpdateScriptRequest(ShardId shardId, UpdateRequest request, GetRes
     private Map<String, Object> executeScript(Script script, Map<String, Object> ctx) {
         try {
             if (scriptService != null) {
-                ExecutableScript.Compiled compiledScript = scriptService.compile(script, ScriptContext.UPDATE);
+                ExecutableScript.Compiled compiledScript = scriptService.compile(script, ExecutableScript.UPDATE_CONTEXT);
                 ExecutableScript executableScript = compiledScript.newInstance(script.getParams());
                 executableScript.setNextVar(ContextFields.CTX, ctx);
                 executableScript.run();

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

@@ -75,7 +75,7 @@ protected void setupInnerHitsContext(QueryShardContext queryShardContext,
         if (innerHitBuilder.getScriptFields() != null) {
             for (SearchSourceBuilder.ScriptField field : innerHitBuilder.getScriptFields()) {
                 SearchScript searchScript = innerHitsContext.getQueryShardContext().getSearchScript(field.script(),
-                    ScriptContext.SEARCH);
+                    SearchScript.CONTEXT);
                 innerHitsContext.scriptFields().add(new org.elasticsearch.search.fetch.subphase.ScriptFieldsContext.ScriptField(
                     field.fieldName(), searchScript, field.ignoreFailure()));
             }

core/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java

@@ -24,6 +24,7 @@
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.mapper.MapperService;
+import org.elasticsearch.script.ExecutableScript;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.ScriptService;
@@ -103,7 +104,7 @@ public long nowInMillis() {
     }
 
     public String getTemplateBytes(Script template) {
-        CompiledTemplate compiledTemplate = scriptService.compileTemplate(template, ScriptContext.EXECUTABLE);
+        CompiledTemplate compiledTemplate = scriptService.compileTemplate(template, ExecutableScript.CONTEXT);
         return compiledTemplate.run(template.getParams());
     }
 }

core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java

@@ -131,7 +131,7 @@ public static ScriptQueryBuilder fromXContent(QueryParseContext parseContext) th
 
     @Override
     protected Query doToQuery(QueryShardContext context) throws IOException {
-        return new ScriptQuery(script, context.getSearchScript(script, ScriptContext.SEARCH));
+        return new ScriptQuery(script, context.getSearchScript(script, SearchScript.CONTEXT));
     }
 
     static class ScriptQuery extends Query {

core/src/main/java/org/elasticsearch/index/query/functionscore/ScriptScoreFunctionBuilder.java

@@ -94,7 +94,7 @@ protected int doHashCode() {
     @Override
     protected ScoreFunction doToFunction(QueryShardContext context) {
         try {
-            SearchScript searchScript = context.getSearchScript(script, ScriptContext.SEARCH);
+            SearchScript searchScript = context.getSearchScript(script, SearchScript.CONTEXT);
             return new ScriptScoreFunction(script, searchScript);
         } catch (Exception e) {
             throw new QueryShardException(context, "script_score: the script could not be loaded", e);

core/src/main/java/org/elasticsearch/ingest/InternalTemplateService.java

@@ -24,6 +24,7 @@
 
 import org.apache.logging.log4j.util.Supplier;
 import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.script.ExecutableScript;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.ScriptService;
@@ -44,7 +45,7 @@ public Template compile(String template) {
         int mustacheEnd = template.indexOf("}}");
         if (mustacheStart != -1 && mustacheEnd != -1 && mustacheStart < mustacheEnd) {
             Script script = new Script(ScriptType.INLINE, "mustache", template, Collections.emptyMap());
-            CompiledTemplate compiledTemplate = scriptService.compileTemplate(script, ScriptContext.INGEST);
+            CompiledTemplate compiledTemplate = scriptService.compileTemplate(script, ExecutableScript.INGEST_CONTEXT);
             return new Template() {
                 @Override
                 public String execute(Map<String, Object> model) {

core/src/main/java/org/elasticsearch/script/ExecutableScript.java

@@ -44,4 +44,11 @@ public interface ExecutableScript {
     interface Compiled {
         ExecutableScript newInstance(Map<String, Object> params);
     }
+
+    ScriptContext<Compiled> CONTEXT = new ScriptContext<>("executable", Compiled.class);
+
+    // TODO: remove these once each has its own script interface
+    ScriptContext<Compiled> AGGS_CONTEXT = new ScriptContext<>("aggs_executable", Compiled.class);
+    ScriptContext<Compiled> UPDATE_CONTEXT = new ScriptContext<>("update", Compiled.class);
+    ScriptContext<Compiled> INGEST_CONTEXT = new ScriptContext<>("ingest", Compiled.class);
 }

core/src/main/java/org/elasticsearch/script/ScriptContext.java

@@ -45,32 +45,6 @@
  */
 public final class ScriptContext<CompiledType> {
 
-    public static final ScriptContext<SearchScript.Compiled> AGGS =
-        new ScriptContext<>("aggs", SearchScript.Compiled.class);
-    public static final ScriptContext<SearchScript.Compiled> SEARCH =
-        new ScriptContext<>("search", SearchScript.Compiled.class);
-    // TODO: remove this once each agg calling scripts has its own context
-    public static final ScriptContext<ExecutableScript.Compiled> AGGS_EXECUTABLE =
-        new ScriptContext<>("aggs_executable", ExecutableScript.Compiled.class);
-    public static final ScriptContext<ExecutableScript.Compiled> UPDATE =
-        new ScriptContext<>("update", ExecutableScript.Compiled.class);
-    public static final ScriptContext<ExecutableScript.Compiled> INGEST =
-        new ScriptContext<>("ingest", ExecutableScript.Compiled.class);
-    public static final ScriptContext<ExecutableScript.Compiled> EXECUTABLE =
-        new ScriptContext<>("executable", ExecutableScript.Compiled.class);
-
-    public static final Map<String, ScriptContext<?>> BUILTINS;
-    static {
-        Map<String, ScriptContext<?>> builtins = new HashMap<>();
-        builtins.put(AGGS.name, AGGS);
-        builtins.put(SEARCH.name, SEARCH);
-        builtins.put(AGGS_EXECUTABLE.name, AGGS_EXECUTABLE);
-        builtins.put(UPDATE.name, UPDATE);
-        builtins.put(INGEST.name, INGEST);
-        builtins.put(EXECUTABLE.name, EXECUTABLE);
-        BUILTINS = Collections.unmodifiableMap(builtins);
-    }
-
     /** A unique identifier for this context. */
     public final String name;
 

core/src/main/java/org/elasticsearch/script/ScriptModule.java

@@ -19,10 +19,14 @@
 
 package org.elasticsearch.script;
 
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.settings.Settings;
@@ -32,11 +36,24 @@
  * Manages building {@link ScriptService}.
  */
 public class ScriptModule {
+
+    public static final Map<String, ScriptContext<?>> CORE_CONTEXTS;
+    static {
+        CORE_CONTEXTS = Stream.of(
+            SearchScript.CONTEXT,
+            SearchScript.AGGS_CONTEXT,
+            ExecutableScript.CONTEXT,
+            ExecutableScript.AGGS_CONTEXT,
+            ExecutableScript.UPDATE_CONTEXT,
+            ExecutableScript.INGEST_CONTEXT
+        ).collect(Collectors.toMap(c -> c.name, Function.identity()));
+    }
+
     private final ScriptService scriptService;
 
     public ScriptModule(Settings settings, List<ScriptPlugin> scriptPlugins) {
         Map<String, ScriptEngine> engines = new HashMap<>();
-        Map<String, ScriptContext<?>> contexts = new HashMap<>(ScriptContext.BUILTINS);
+        Map<String, ScriptContext<?>> contexts = new HashMap<>(CORE_CONTEXTS);
         for (ScriptPlugin plugin : scriptPlugins) {
             for (ScriptContext context : plugin.getContexts()) {
                 ScriptContext oldContext = contexts.put(context.name, context);

core/src/main/java/org/elasticsearch/script/ScriptService.java

@@ -264,7 +264,7 @@ public <CompiledType> CompiledType compile(Script script, ScriptContext<Compiled
         // TODO: fix this through some API or something, that's wrong
         // special exception to prevent expressions from compiling as update or mapping scripts
         boolean expression = "expression".equals(script.getLang());
-        boolean notSupported = context.name.equals(ScriptContext.UPDATE.name);
+        boolean notSupported = context.name.equals(ExecutableScript.UPDATE_CONTEXT.name);
         if (expression && notSupported) {
             throw new UnsupportedOperationException("scripts of type [" + script.getType() + "]," +
                 " operation [" + context.name + "] and lang [" + lang + "] are not supported");
@@ -433,7 +433,7 @@ public void putStoredScript(ClusterService clusterService, PutStoredScriptReques
                     "cannot put [" + ScriptType.STORED + "] script, no script contexts are enabled");
             } else {
                 // TODO: executable context here is just a placeholder, replace with optional context name passed into PUT stored script req
-                Object compiled = scriptEngine.compile(request.id(), source.getCode(), ScriptContext.EXECUTABLE, Collections.emptyMap());
+                Object compiled = scriptEngine.compile(request.id(), source.getCode(), ExecutableScript.CONTEXT, Collections.emptyMap());
 
                 if (compiled == null) {
                     throw new IllegalArgumentException("failed to parse/compile stored script [" + request.id() + "]" +

core/src/main/java/org/elasticsearch/script/SearchScript.java

@@ -41,4 +41,8 @@ public interface SearchScript {
     interface Compiled {
         SearchScript newInstance(Map<String, Object> params, SearchLookup lookup);
     }
+
+    ScriptContext<Compiled> CONTEXT = new ScriptContext<>("search", Compiled.class);
+    // TODO: remove aggs context when it has its own interface
+    ScriptContext<SearchScript.Compiled> AGGS_CONTEXT = new ScriptContext<>("aggs", Compiled.class);
 }

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

@@ -685,7 +685,7 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
         }
         if (source.scriptFields() != null) {
             for (org.elasticsearch.search.builder.SearchSourceBuilder.ScriptField field : source.scriptFields()) {
-                SearchScript.Compiled compiled = scriptService.compile(field.script(), ScriptContext.SEARCH);
+                SearchScript.Compiled compiled = scriptService.compile(field.script(), SearchScript.CONTEXT);
                 SearchScript searchScript = compiled.newInstance(field.script().getParams(), context.lookup());
                 context.scriptFields().add(new ScriptField(field.fieldName(), searchScript, field.ignoreFailure()));
             }

core/src/main/java/org/elasticsearch/search/aggregations/bucket/significant/heuristics/ScriptHeuristic.java

@@ -92,14 +92,14 @@ public void writeTo(StreamOutput out) throws IOException {
 
     @Override
     public SignificanceHeuristic rewrite(InternalAggregation.ReduceContext context) {
-        ExecutableScript.Compiled compiled = context.scriptService().compile(script, ScriptContext.AGGS_EXECUTABLE);
+        ExecutableScript.Compiled compiled = context.scriptService().compile(script, ExecutableScript.AGGS_CONTEXT);
         return new ExecutableScriptHeuristic(script, compiled.newInstance(script.getParams()));
     }
 
     @Override
     public SignificanceHeuristic rewrite(SearchContext context) {
         return new ExecutableScriptHeuristic(script,
-            context.getQueryShardContext().getExecutableScript(script, ScriptContext.AGGS_EXECUTABLE));
+            context.getQueryShardContext().getExecutableScript(script, ExecutableScript.AGGS_CONTEXT));
     }
 
 

core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/InternalScriptedMetric.java

@@ -96,7 +96,7 @@ public InternalAggregation doReduce(List<InternalAggregation> aggregations, Redu
                 vars.putAll(firstAggregation.reduceScript.getParams());
             }
             ExecutableScript.Compiled compiled = reduceContext.scriptService().compile(
-                firstAggregation.reduceScript, ScriptContext.AGGS_EXECUTABLE);
+                firstAggregation.reduceScript, ExecutableScript.AGGS_CONTEXT);
             ExecutableScript script = compiled.newInstance(vars);
             aggregation = Collections.singletonList(script.run());
         } else if (reduceContext.isFinalReduce())  {

core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregationBuilder.java

@@ -186,14 +186,15 @@ protected ScriptedMetricAggregatorFactory doBuild(SearchContext context, Aggrega
         QueryShardContext queryShardContext = context.getQueryShardContext();
         Function<Map<String, Object>, ExecutableScript> executableInitScript;
         if (initScript != null) {
-            executableInitScript = queryShardContext.getLazyExecutableScript(initScript, ScriptContext.AGGS_EXECUTABLE);
+            executableInitScript = queryShardContext.getLazyExecutableScript(initScript, ExecutableScript.AGGS_CONTEXT);
         } else {
             executableInitScript = (p) -> null;
         }
-        Function<Map<String, Object>, SearchScript> searchMapScript = queryShardContext.getLazySearchScript(mapScript, ScriptContext.AGGS);
+        Function<Map<String, Object>, SearchScript> searchMapScript =
+            queryShardContext.getLazySearchScript(mapScript, SearchScript.AGGS_CONTEXT);
         Function<Map<String, Object>, ExecutableScript> executableCombineScript;
         if (combineScript != null) {
-            executableCombineScript = queryShardContext.getLazyExecutableScript(combineScript, ScriptContext.AGGS_EXECUTABLE);
+            executableCombineScript = queryShardContext.getLazyExecutableScript(combineScript, ExecutableScript.AGGS_CONTEXT);
         } else {
             executableCombineScript = (p) -> null;
         }

core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregationBuilder.java

@@ -534,7 +534,7 @@ protected TopHitsAggregatorFactory doBuild(SearchContext context, AggregatorFact
         if (scriptFields != null) {
             for (ScriptField field : scriptFields) {
                 SearchScript searchScript = context.getQueryShardContext().getSearchScript(field.script(),
-                    ScriptContext.SEARCH);
+                    SearchScript.CONTEXT);
                 fields.add(new org.elasticsearch.search.fetch.subphase.ScriptFieldsContext.ScriptField(
                     field.fieldName(), searchScript, field.ignoreFailure()));
             }

core/src/main/java/org/elasticsearch/search/aggregations/pipeline/bucketscript/BucketScriptPipelineAggregator.java

@@ -90,7 +90,7 @@ public InternalAggregation reduce(InternalAggregation aggregation, ReduceContext
                 (InternalMultiBucketAggregation<InternalMultiBucketAggregation, InternalMultiBucketAggregation.InternalBucket>) aggregation;
         List<? extends InternalMultiBucketAggregation.InternalBucket> buckets = originalAgg.getBuckets();
 
-        ExecutableScript.Compiled compiledScript = reduceContext.scriptService().compile(script, ScriptContext.AGGS_EXECUTABLE);
+        ExecutableScript.Compiled compiledScript = reduceContext.scriptService().compile(script, ExecutableScript.AGGS_CONTEXT);
         List<InternalMultiBucketAggregation.InternalBucket> newBuckets = new ArrayList<>();
         for (InternalMultiBucketAggregation.InternalBucket bucket : buckets) {
             Map<String, Object> vars = new HashMap<>();

core/src/main/java/org/elasticsearch/search/aggregations/pipeline/bucketselector/BucketSelectorPipelineAggregator.java

@@ -83,7 +83,7 @@ public InternalAggregation reduce(InternalAggregation aggregation, ReduceContext
                 (InternalMultiBucketAggregation<InternalMultiBucketAggregation, InternalMultiBucketAggregation.InternalBucket>) aggregation;
         List<? extends InternalMultiBucketAggregation.InternalBucket> buckets = originalAgg.getBuckets();
 
-        ExecutableScript.Compiled compiledScript = reduceContext.scriptService().compile(script, ScriptContext.AGGS_EXECUTABLE);
+        ExecutableScript.Compiled compiledScript = reduceContext.scriptService().compile(script, ExecutableScript.AGGS_CONTEXT);
         List<InternalMultiBucketAggregation.InternalBucket> newBuckets = new ArrayList<>();
         for (InternalMultiBucketAggregation.InternalBucket bucket : buckets) {
             Map<String, Object> vars = new HashMap<>();

core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceConfig.java

@@ -120,7 +120,7 @@ private static SearchScript createScript(Script script, QueryShardContext contex
         if (script == null) {
             return null;
         } else {
-            return context.getSearchScript(script, ScriptContext.AGGS);
+            return context.getSearchScript(script, SearchScript.AGGS_CONTEXT);
         }
     }
 

core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java

@@ -242,7 +242,7 @@ public static ScriptSortBuilder fromXContent(QueryParseContext context, String e
 
     @Override
     public SortFieldAndFormat build(QueryShardContext context) throws IOException {
-        final SearchScript searchScript = context.getSearchScript(script, ScriptContext.SEARCH);
+        final SearchScript searchScript = context.getSearchScript(script, SearchScript.CONTEXT);
 
         MultiValueMode valueMode = null;
         if (sortMode != null) {

core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggestionBuilder.java

@@ -631,7 +631,7 @@ public SuggestionContext build(QueryShardContext context) throws IOException {
 
         if (this.collateQuery != null) {
             Function<Map<String, Object>, ExecutableScript> compiledScript = context.getLazyExecutableScript(this.collateQuery,
-                ScriptContext.EXECUTABLE);
+                ExecutableScript.CONTEXT);
             suggestionContext.setCollateQueryScript(compiledScript);
             if (this.collateParams != null) {
                 suggestionContext.setCollateScriptParams(this.collateParams);

core/src/test/java/org/elasticsearch/action/update/UpdateRequestTests.java

@@ -42,8 +42,8 @@
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.script.MockScriptEngine;
 import org.elasticsearch.script.Script;
-import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.ScriptEngine;
+import org.elasticsearch.script.ScriptModule;
 import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.script.ScriptType;
 import org.elasticsearch.test.ESTestCase;
@@ -136,7 +136,7 @@ public void setUp() throws Exception {
         scripts.put("return", vars -> null);
         final MockScriptEngine engine = new MockScriptEngine("mock", scripts);
         Map<String, ScriptEngine> engines = Collections.singletonMap(engine.getType(), engine);
-        ScriptService scriptService = new ScriptService(baseSettings, engines, ScriptContext.BUILTINS);
+        ScriptService scriptService = new ScriptService(baseSettings, engines, ScriptModule.CORE_CONTEXTS);
         final Settings settings = settings(Version.CURRENT).build();
 
         updateHelper = new UpdateHelper(settings, scriptService);