Add more script contexts

plugins/examples/script-expert-scoring/src/main/java/org/elasticsearch/example/expertscript/ExpertScriptPlugin.java

@@ -54,7 +54,7 @@ public String getType() {
 
         @Override
         public <T> T compile(String scriptName, String scriptSource, ScriptContext<T> context, Map<String, String> params) {
-            if (context.equals(SearchScript.CONTEXT) == false) {
+            if (context.equals(SearchScript.SCRIPT_SCORE_CONTEXT) == false) {
                 throw new IllegalArgumentException(getType() + " scripts cannot be used for context [" + context.name + "]");
             }
             // we use the script "source" as the script identifier

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

@@ -249,7 +249,8 @@ protected Query doToQuery(QueryShardContext context) throws IOException {
             IndexNumericFieldData fieldData = context.getForField(msmFieldType);
             longValuesSource = new FieldValuesSource(fieldData);
         } else if (minimumShouldMatchScript != null) {
-            SearchScript.Factory factory = context.getScriptService().compile(minimumShouldMatchScript, SearchScript.CONTEXT);
+            SearchScript.Factory factory = context.getScriptService().compile(minimumShouldMatchScript,
+                SearchScript.TERMS_SET_QUERY_CONTEXT);
             Map<String, Object> params = new HashMap<>();
             params.putAll(minimumShouldMatchScript.getParams());
             params.put("num_terms", queries.size());

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

@@ -92,7 +92,7 @@ protected int doHashCode() {
     @Override
     protected ScoreFunction doToFunction(QueryShardContext context) {
         try {
-            SearchScript.Factory factory = context.getScriptService().compile(script, SearchScript.CONTEXT);
+            SearchScript.Factory factory = context.getScriptService().compile(script, SearchScript.SCRIPT_SCORE_CONTEXT);
             SearchScript.LeafFactory searchScript = factory.newFactory(script.getParams(), context.lookup());
             return new ScriptScoreFunction(script, searchScript);
         } catch (Exception e) {

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

@@ -42,6 +42,9 @@ public class ScriptModule {
         CORE_CONTEXTS = Stream.of(
             SearchScript.CONTEXT,
             SearchScript.AGGS_CONTEXT,
+            SearchScript.SCRIPT_SCORE_CONTEXT,
+            SearchScript.SCRIPT_SORT_CONTEXT,
+            SearchScript.TERMS_SET_QUERY_CONTEXT,
             ExecutableScript.CONTEXT,
             ExecutableScript.AGGS_CONTEXT,
             ExecutableScript.UPDATE_CONTEXT,

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

@@ -158,6 +158,12 @@ public interface Factory {
 
     /** The context used to compile {@link SearchScript} factories. */
     public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("search", Factory.class);
-    // TODO: remove aggs context when it has its own interface
+    // TODO: remove these contexts when it has its own interface
     public static final ScriptContext<Factory> AGGS_CONTEXT = new ScriptContext<>("aggs", Factory.class);
-}
+    // Can return a double. (For ScriptSortType#NUMBER only, for ScriptSortType#STRING normal CONTEXT should be used)
+    public static final ScriptContext<Factory> SCRIPT_SORT_CONTEXT = new ScriptContext<>("sort", Factory.class);
+    // Can return a float
+    public static final ScriptContext<Factory> SCRIPT_SCORE_CONTEXT = new ScriptContext<>("score", Factory.class);
+    // Can return a long
+    public static final ScriptContext<Factory> TERMS_SET_QUERY_CONTEXT = new ScriptContext<>("terms_set", Factory.class);
+}

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

@@ -305,7 +305,7 @@ public static ScriptSortBuilder fromXContent(XContentParser parser, String eleme
 
     @Override
     public SortFieldAndFormat build(QueryShardContext context) throws IOException {
-        final SearchScript.Factory factory = context.getScriptService().compile(script, SearchScript.CONTEXT);
+        final SearchScript.Factory factory = context.getScriptService().compile(script, SearchScript.SCRIPT_SORT_CONTEXT);
         final SearchScript.LeafFactory searchScript = factory.newFactory(script.getParams(), context.lookup());
 
         MultiValueMode valueMode = null;

server/src/test/java/org/elasticsearch/search/functionscore/ExplainableScriptIT.java

@@ -76,7 +76,7 @@ public String getType() {
                 @Override
                 public <T> T compile(String scriptName, String scriptSource, ScriptContext<T> context, Map<String, String> params) {
                     assert scriptSource.equals("explainable_script");
-                    assert context == SearchScript.CONTEXT;
+                    assert context == SearchScript.SCRIPT_SCORE_CONTEXT;
                     SearchScript.Factory factory = (p, lookup) -> new SearchScript.LeafFactory() {
                         @Override
                         public SearchScript newInstance(LeafReaderContext context) throws IOException {