Unify supported runtime fields script contexts (elastic#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.

Author: javanna

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

@@ -37,17 +37,11 @@
 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

@@ -72,6 +72,7 @@
 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;
@@ -115,23 +116,19 @@ public static class Request extends SingleShardRequest<Request> implements ToXCo
             PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(), ContextSetup::parse, CONTEXT_SETUP_FIELD);
         }
 
-        static final Map<String, ScriptContext<?>> SUPPORTED_CONTEXTS;
-
-        static {
-            Map<String, ScriptContext<?>> supportedContexts = new HashMap<>();
-            supportedContexts.put("painless_test", PainlessTestScript.CONTEXT);
-            supportedContexts.put("filter", FilterScript.CONTEXT);
-            supportedContexts.put("score", ScoreScript.CONTEXT);
-            supportedContexts.put("boolean_field", BooleanFieldScript.CONTEXT);
-            supportedContexts.put("date_field", DateFieldScript.CONTEXT);
-            supportedContexts.put("double_field", DoubleFieldScript.CONTEXT);
-            supportedContexts.put("geo_point_field", GeoPointFieldScript.CONTEXT);
-            supportedContexts.put("ip_field", IpFieldScript.CONTEXT);
-            supportedContexts.put("long_field", LongFieldScript.CONTEXT);
-            supportedContexts.put("string_field", StringFieldScript.CONTEXT);
-            SUPPORTED_CONTEXTS = Collections.unmodifiableMap(supportedContexts);
+        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);
             if (scriptContext == null) {
@@ -561,15 +558,15 @@ static Response innerShardOperation(Request request, ScriptService scriptService
                 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());
                 }, indexService);
             } 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);
@@ -579,7 +576,7 @@ static Response innerShardOperation(Request request, ScriptService scriptService
                 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());
@@ -588,7 +585,7 @@ static Response innerShardOperation(Request request, ScriptService scriptService
                 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());
@@ -597,7 +594,7 @@ static Response innerShardOperation(Request request, ScriptService scriptService
                 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());
@@ -606,7 +603,7 @@ static Response innerShardOperation(Request request, ScriptService scriptService
                 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());
@@ -615,7 +612,7 @@ static Response innerShardOperation(Request request, ScriptService scriptService
                 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/lang-painless/src/main/resources/org/elasticsearch/painless/spi/org.elasticsearch.script.boolean.txt renamed to modules/lang-painless/src/main/resources/org/elasticsearch/painless/spi/org.elasticsearch.script.boolean_field.txt

@@ -14,15 +14,11 @@
 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 RuntimeFieldsPainlessExtension(RuntimeFieldsCommonPlugin plugin) {
 
     @Override
     public Map<ScriptContext<?>, List<Whitelist>> getContextWhitelists() {
-        return org.elasticsearch.common.collect.Map.ofEntries(
-            org.elasticsearch.common.collect.Map.entry(BooleanFieldScript.CONTEXT, whitelists),
-            org.elasticsearch.common.collect.Map.entry(DateFieldScript.CONTEXT, whitelists),
-            org.elasticsearch.common.collect.Map.entry(DoubleFieldScript.CONTEXT, whitelists),
-            org.elasticsearch.common.collect.Map.entry(GeoPointFieldScript.CONTEXT, whitelists),
-            org.elasticsearch.common.collect.Map.entry(IpFieldScript.CONTEXT, whitelists),
-            org.elasticsearch.common.collect.Map.entry(LongFieldScript.CONTEXT, whitelists),
-            org.elasticsearch.common.collect.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);
     }
 }

modules/lang-painless/src/main/resources/org/elasticsearch/painless/spi/org.elasticsearch.script.date.txt renamed to modules/lang-painless/src/main/resources/org/elasticsearch/painless/spi/org.elasticsearch.script.date_field.txt

@@ -18,6 +18,7 @@
 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,13 @@
  */
 public class ScriptModule {
 
+    public static final Set<ScriptContext<?>> RUNTIME_FIELDS_CONTEXTS = org.elasticsearch.common.collect.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 +57,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;