Unify supported runtime fields script contexts (#71833)

There's a few places where we need to access all of the supported runtime fields script contexts. Up until now we have listed them in all those places, but a better way would be to have them listed in one place and access that same list from all consumers. This is what this commit introduces.

Along with the introduction of runtime fields contexts in ScriptModule, we rename the whitelist files so that they contain their corresponding context name to simplify looking them up.

modules/lang-painless/src/main/java/org/elasticsearch/painless/PainlessPlugin.java

@@ -37,17 +37,11 @@ import org.elasticsearch.plugins.ScriptPlugin;
 import org.elasticsearch.repositories.RepositoriesService;
 import org.elasticsearch.rest.RestController;
 import org.elasticsearch.rest.RestHandler;
-import org.elasticsearch.script.BooleanFieldScript;
-import org.elasticsearch.script.DateFieldScript;
-import org.elasticsearch.script.DoubleFieldScript;
-import org.elasticsearch.script.GeoPointFieldScript;
-import org.elasticsearch.script.IpFieldScript;
-import org.elasticsearch.script.LongFieldScript;
-import org.elasticsearch.script.StringFieldScript;
 import org.elasticsearch.script.IngestScript;
 import org.elasticsearch.script.ScoreScript;
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.ScriptEngine;
+import org.elasticsearch.script.ScriptModule;
 import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.search.aggregations.pipeline.MovingFunctionScript;
 import org.elasticsearch.threadpool.ThreadPool;
@@ -96,13 +90,10 @@ public final class PainlessPlugin extends Plugin implements ScriptPlugin, Extens
         map.put(IngestScript.CONTEXT, ingest);
 
         // Functions available to runtime fields
-        map.put(BooleanFieldScript.CONTEXT, getRuntimeFieldWhitelist("boolean"));
-        map.put(DateFieldScript.CONTEXT, getRuntimeFieldWhitelist("date"));
-        map.put(DoubleFieldScript.CONTEXT, getRuntimeFieldWhitelist("double"));
-        map.put(LongFieldScript.CONTEXT, getRuntimeFieldWhitelist("long"));
-        map.put(StringFieldScript.CONTEXT, getRuntimeFieldWhitelist("string"));
-        map.put(GeoPointFieldScript.CONTEXT, getRuntimeFieldWhitelist("geopoint"));
-        map.put(IpFieldScript.CONTEXT, getRuntimeFieldWhitelist("ip"));
+
+        for (ScriptContext<?> scriptContext : ScriptModule.RUNTIME_FIELDS_CONTEXTS) {
+            map.put(scriptContext, getRuntimeFieldWhitelist(scriptContext.name));
+        }
 
         // Execute context gets everything
         List<Whitelist> test = new ArrayList<>(Whitelist.BASE_WHITELISTS);
@@ -115,10 +106,10 @@ public final class PainlessPlugin extends Plugin implements ScriptPlugin, Extens
         whitelists = map;
     }
 
-    private static List<Whitelist> getRuntimeFieldWhitelist(String fieldType) {
+    private static List<Whitelist> getRuntimeFieldWhitelist(String contextName) {
         List<Whitelist> scriptField = new ArrayList<>(Whitelist.BASE_WHITELISTS);
         Whitelist whitelist = WhitelistLoader.loadFromResourceFiles(Whitelist.class,
-            "org.elasticsearch.script." + fieldType + ".txt");
+            "org.elasticsearch.script." + contextName + ".txt");
         scriptField.add(whitelist);
         return scriptField;
     }

modules/lang-painless/src/main/java/org/elasticsearch/painless/action/PainlessExecuteAction.java

@@ -71,6 +71,7 @@ import org.elasticsearch.script.LongFieldScript;
 import org.elasticsearch.script.ScoreScript;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptContext;
+import org.elasticsearch.script.ScriptModule;
 import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.script.ScriptType;
 import org.elasticsearch.script.StringFieldScript;
@@ -78,6 +79,8 @@ import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 
 import java.io.IOException;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -110,18 +113,18 @@ public class PainlessExecuteAction extends ActionType<PainlessExecuteAction.Resp
             PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(), ContextSetup::parse, CONTEXT_SETUP_FIELD);
         }
 
-        static final Map<String, ScriptContext<?>> SUPPORTED_CONTEXTS = Map.of(
-                "painless_test", PainlessTestScript.CONTEXT,
-                "filter", FilterScript.CONTEXT,
-                "score", ScoreScript.CONTEXT,
-                "boolean_field", BooleanFieldScript.CONTEXT,
-                "date_field", DateFieldScript.CONTEXT,
-                "double_field", DoubleFieldScript.CONTEXT,
-                "geo_point_field", GeoPointFieldScript.CONTEXT,
-                "ip_field", IpFieldScript.CONTEXT,
-                "long_field", LongFieldScript.CONTEXT,
-                "string_field", StringFieldScript.CONTEXT)
-        ;
+        private static Map<String, ScriptContext<?>> getSupportedContexts() {
+            Map<String, ScriptContext<?>> contexts = new HashMap<>();
+            contexts.put(PainlessTestScript.CONTEXT.name, PainlessTestScript.CONTEXT);
+            contexts.put(FilterScript.CONTEXT.name, FilterScript.CONTEXT);
+            contexts.put(ScoreScript.CONTEXT.name, ScoreScript.CONTEXT);
+            for (ScriptContext<?> runtimeFieldsContext : ScriptModule.RUNTIME_FIELDS_CONTEXTS) {
+                contexts.put(runtimeFieldsContext.name, runtimeFieldsContext);
+            }
+            return Collections.unmodifiableMap(contexts);
+        }
+
+        static final Map<String, ScriptContext<?>> SUPPORTED_CONTEXTS = getSupportedContexts();
 
         static ScriptContext<?> fromScriptContextName(String name) {
             ScriptContext<?> scriptContext = SUPPORTED_CONTEXTS.get(name);
@@ -541,7 +544,7 @@ public class PainlessExecuteAction extends ActionType<PainlessExecuteAction.Resp
                 return prepareRamIndex(request, (context, leafReaderContext) -> {
                     BooleanFieldScript.Factory factory = scriptService.compile(request.script, BooleanFieldScript.CONTEXT);
                     BooleanFieldScript.LeafFactory leafFactory =
-                            factory.newFactory("boolean_field", request.getScript().getParams(), context.lookup());
+                            factory.newFactory(BooleanFieldScript.CONTEXT.name, request.getScript().getParams(), context.lookup());
                     BooleanFieldScript booleanFieldScript = leafFactory.newInstance(leafReaderContext);
                     booleanFieldScript.runForDoc(0);
                     return new Response(booleanFieldScript.asDocValues());
@@ -549,7 +552,7 @@ public class PainlessExecuteAction extends ActionType<PainlessExecuteAction.Resp
             } else if (scriptContext == DateFieldScript.CONTEXT) {
                 return prepareRamIndex(request, (context, leafReaderContext) -> {
                     DateFieldScript.Factory factory = scriptService.compile(request.script, DateFieldScript.CONTEXT);
-                    DateFieldScript.LeafFactory leafFactory = factory.newFactory("date_field",
+                    DateFieldScript.LeafFactory leafFactory = factory.newFactory(DateFieldScript.CONTEXT.name,
                             request.getScript().getParams(), context.lookup(), DateFieldMapper.DEFAULT_DATE_TIME_FORMATTER);
                     DateFieldScript dateFieldScript = leafFactory.newInstance(leafReaderContext);
                     dateFieldScript.runForDoc(0);
@@ -559,7 +562,7 @@ public class PainlessExecuteAction extends ActionType<PainlessExecuteAction.Resp
                 return prepareRamIndex(request, (context, leafReaderContext) -> {
                     DoubleFieldScript.Factory factory = scriptService.compile(request.script, DoubleFieldScript.CONTEXT);
                     DoubleFieldScript.LeafFactory leafFactory =
-                            factory.newFactory("double_field", request.getScript().getParams(), context.lookup());
+                            factory.newFactory(DoubleFieldScript.CONTEXT.name, request.getScript().getParams(), context.lookup());
                     DoubleFieldScript doubleFieldScript = leafFactory.newInstance(leafReaderContext);
                     doubleFieldScript.runForDoc(0);
                     return new Response(doubleFieldScript.asDocValues());
@@ -568,7 +571,7 @@ public class PainlessExecuteAction extends ActionType<PainlessExecuteAction.Resp
                 return prepareRamIndex(request, (context, leafReaderContext) -> {
                     GeoPointFieldScript.Factory factory = scriptService.compile(request.script, GeoPointFieldScript.CONTEXT);
                     GeoPointFieldScript.LeafFactory leafFactory =
-                            factory.newFactory("geo_point_field", request.getScript().getParams(), context.lookup());
+                            factory.newFactory(GeoPointFieldScript.CONTEXT.name, request.getScript().getParams(), context.lookup());
                     GeoPointFieldScript geoPointFieldScript = leafFactory.newInstance(leafReaderContext);
                     geoPointFieldScript.runForDoc(0);
                     return new Response(geoPointFieldScript.asDocValues());
@@ -577,7 +580,7 @@ public class PainlessExecuteAction extends ActionType<PainlessExecuteAction.Resp
                 return prepareRamIndex(request, (context, leafReaderContext) -> {
                     IpFieldScript.Factory factory = scriptService.compile(request.script, IpFieldScript.CONTEXT);
                     IpFieldScript.LeafFactory leafFactory =
-                            factory.newFactory("ip_field", request.getScript().getParams(), context.lookup());
+                            factory.newFactory(IpFieldScript.CONTEXT.name, request.getScript().getParams(), context.lookup());
                     IpFieldScript ipFieldScript = leafFactory.newInstance(leafReaderContext);
                     ipFieldScript.runForDoc(0);
                     return new Response(ipFieldScript.asDocValues());
@@ -586,7 +589,7 @@ public class PainlessExecuteAction extends ActionType<PainlessExecuteAction.Resp
                 return prepareRamIndex(request, (context, leafReaderContext) -> {
                     LongFieldScript.Factory factory = scriptService.compile(request.script, LongFieldScript.CONTEXT);
                     LongFieldScript.LeafFactory leafFactory =
-                            factory.newFactory("long_field", request.getScript().getParams(), context.lookup());
+                            factory.newFactory(LongFieldScript.CONTEXT.name, request.getScript().getParams(), context.lookup());
                     LongFieldScript longFieldScript = leafFactory.newInstance(leafReaderContext);
                     longFieldScript.runForDoc(0);
                     return new Response(longFieldScript.asDocValues());
@@ -595,7 +598,7 @@ public class PainlessExecuteAction extends ActionType<PainlessExecuteAction.Resp
                 return prepareRamIndex(request, (context, leafReaderContext) -> {
                     StringFieldScript.Factory factory = scriptService.compile(request.script, StringFieldScript.CONTEXT);
                     StringFieldScript.LeafFactory leafFactory =
-                            factory.newFactory("string_field", request.getScript().getParams(), context.lookup());
+                            factory.newFactory(StringFieldScript.CONTEXT.name, request.getScript().getParams(), context.lookup());
                     StringFieldScript stringFieldScript = leafFactory.newInstance(leafReaderContext);
                     stringFieldScript.resultsForDoc(0);
                     return new Response(stringFieldScript.asDocValues());

modules/runtime-fields-common/src/main/java/org/elasticsearch/runtimefields/RuntimeFieldsPainlessExtension.java

@@ -14,15 +14,11 @@ import org.elasticsearch.painless.spi.WhitelistInstanceBinding;
 import org.elasticsearch.painless.spi.WhitelistLoader;
 import org.elasticsearch.painless.spi.annotation.CompileTimeOnlyAnnotation;
 import org.elasticsearch.script.AbstractFieldScript;
-import org.elasticsearch.script.BooleanFieldScript;
-import org.elasticsearch.script.DateFieldScript;
-import org.elasticsearch.script.DoubleFieldScript;
-import org.elasticsearch.script.GeoPointFieldScript;
-import org.elasticsearch.script.IpFieldScript;
-import org.elasticsearch.script.LongFieldScript;
-import org.elasticsearch.script.StringFieldScript;
 import org.elasticsearch.script.ScriptContext;
+import org.elasticsearch.script.ScriptModule;
 
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -55,14 +51,10 @@ public class RuntimeFieldsPainlessExtension implements PainlessExtension {
 
     @Override
     public Map<ScriptContext<?>, List<Whitelist>> getContextWhitelists() {
-        return Map.ofEntries(
-            Map.entry(BooleanFieldScript.CONTEXT, whitelists),
-            Map.entry(DateFieldScript.CONTEXT, whitelists),
-            Map.entry(DoubleFieldScript.CONTEXT, whitelists),
-            Map.entry(GeoPointFieldScript.CONTEXT, whitelists),
-            Map.entry(IpFieldScript.CONTEXT, whitelists),
-            Map.entry(LongFieldScript.CONTEXT, whitelists),
-            Map.entry(StringFieldScript.CONTEXT, whitelists)
-        );
+        Map<ScriptContext<?>, List<Whitelist>> whiteLists = new HashMap<>();
+        for (ScriptContext<?> scriptContext : ScriptModule.RUNTIME_FIELDS_CONTEXTS) {
+            whiteLists.put(scriptContext, whitelists);
+        }
+        return Collections.unmodifiableMap(whiteLists);
     }
 }

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

@@ -18,6 +18,7 @@ import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -27,9 +28,12 @@ import java.util.stream.Stream;
  */
 public class ScriptModule {
 
+    public static final Set<ScriptContext<?>> RUNTIME_FIELDS_CONTEXTS = Set.of(BooleanFieldScript.CONTEXT, DateFieldScript.CONTEXT,
+        DoubleFieldScript.CONTEXT, LongFieldScript.CONTEXT, StringFieldScript.CONTEXT, GeoPointFieldScript.CONTEXT, IpFieldScript.CONTEXT);
+
     public static final Map<String, ScriptContext<?>> CORE_CONTEXTS;
     static {
-        CORE_CONTEXTS = Stream.of(
+        CORE_CONTEXTS = Stream.concat(Stream.of(
             FieldScript.CONTEXT,
             AggregationScript.CONTEXT,
             ScoreScript.CONTEXT,
@@ -52,15 +56,8 @@ public class ScriptModule {
             ScriptedMetricAggContexts.MapScript.CONTEXT,
             ScriptedMetricAggContexts.CombineScript.CONTEXT,
             ScriptedMetricAggContexts.ReduceScript.CONTEXT,
-            IntervalFilterScript.CONTEXT,
-            BooleanFieldScript.CONTEXT,
-            DateFieldScript.CONTEXT,
-            DoubleFieldScript.CONTEXT,
-            GeoPointFieldScript.CONTEXT,
-            IpFieldScript.CONTEXT,
-            LongFieldScript.CONTEXT,
-            StringFieldScript.CONTEXT
-        ).collect(Collectors.toMap(c -> c.name, Function.identity()));
+            IntervalFilterScript.CONTEXT
+        ), RUNTIME_FIELDS_CONTEXTS.stream()).collect(Collectors.toMap(c -> c.name, Function.identity()));
     }
 
     public final Map<String, ScriptEngine> engines;