Scripted fields: rework script contexts

x-pack/plugin/runtime-fields/build.gradle

@@ -6,7 +6,7 @@ esplugin {
   name 'x-pack-runtime-fields'
   description 'A module which adds support for runtime fields'
   classname 'org.elasticsearch.xpack.runtimefields.RuntimeFields'
-  extendedPlugins = ['x-pack-core']
+  extendedPlugins = ['x-pack-core', 'lang-painless']
 }
 archivesBaseName = 'x-pack-runtime-fields'
 

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/AbstractScriptFieldsScript.java

@@ -35,29 +35,32 @@ public AbstractScriptFieldsScript(Map<String, Object> params, SourceLookup sourc
     /**
      * Set the document to run the script against.
      */
-    public void setDocument(int docId) {
+    public final void setDocument(int docId) {
         source.setSegmentAndDocument(ctx, docId);
         fieldData.setDocument(docId);
     }
 
     /**
      * Expose the {@code params} of the script to the script itself.
      */
-    public Map<String, Object> getParams() {
+    public final Map<String, Object> getParams() {
         return params;
     }
 
     /**
      * Expose the {@code _source} to the script.
      */
-    public Map<String, Object> getSource() {
+    public final Map<String, Object> getSource() {
         return source;
     }
 
     /**
      * Expose field data to the script as {@code doc}.
      */
-    public Map<String, ScriptDocValues<?>> getDoc() {
+    public final Map<String, ScriptDocValues<?>> getDoc() {
         return fieldData;
     }
+
+    public abstract void execute();
+
 }

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/LongScriptFieldScript.java

@@ -7,28 +7,55 @@
 package org.elasticsearch.xpack.runtimefields;
 
 import org.apache.lucene.index.LeafReaderContext;
+import org.elasticsearch.painless.spi.Whitelist;
+import org.elasticsearch.painless.spi.WhitelistLoader;
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.ScriptFactory;
 import org.elasticsearch.search.lookup.DocLookup;
 import org.elasticsearch.search.lookup.SourceLookup;
 
 import java.io.IOException;
+import java.util.List;
 import java.util.Map;
+import java.util.function.LongConsumer;
 
 public abstract class LongScriptFieldScript extends AbstractScriptFieldsScript {
-    public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("long_script_field", Factory.class);
+    static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("long_script_field", Factory.class);
+    static List<Whitelist> whitelist() {
+        return List.of(WhitelistLoader.loadFromResourceFiles(RuntimeFieldsPainlessExtension.class, "long_whitelist.txt"));
+    }
+
     public static final String[] PARAMETERS = {};
 
     public interface Factory extends ScriptFactory {
         LeafFactory newFactory(Map<String, Object> params, SourceLookup source, DocLookup fieldData);
     }
     public static interface LeafFactory {
-        LongScriptFieldScript newInstance(LeafReaderContext ctx) throws IOException;
+        LongScriptFieldScript newInstance(LeafReaderContext ctx, LongConsumer sync) throws IOException;
     }
 
-    public LongScriptFieldScript(Map<String, Object> params, SourceLookup source, DocLookup fieldData, LeafReaderContext ctx) {
+    private final LongConsumer sync;
+
+    public LongScriptFieldScript(
+        Map<String, Object> params,
+        SourceLookup source,
+        DocLookup fieldData,
+        LeafReaderContext ctx,
+        LongConsumer sync
+    ) {
         super(params, source, fieldData, ctx);
+        this.sync = sync;
     }
 
-    public abstract long execute();
+    public static class Value {
+        private final LongScriptFieldScript script;
+
+        public Value(LongScriptFieldScript script) {
+            this.script = script;
+        }
+
+        public void value(long v) {
+            script.sync.accept(v);
+        }
+    }
 }

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/RuntimeFields.java

@@ -26,9 +26,8 @@ public Map<String, Mapper.TypeParser> getMappers() {
     @Override
     public List<ScriptContext<?>> getContexts() {
         return List.of(
-            LongArrayScriptFieldScript.CONTEXT,
             LongScriptFieldScript.CONTEXT,
-            StringScriptFieldsScript.CONTEXT
+            StringScriptFieldScript.CONTEXT
         );
     }
 }

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/RuntimeFieldsPainlessExtension.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+package org.elasticsearch.xpack.runtimefields;
+
+import org.elasticsearch.painless.spi.PainlessExtension;
+import org.elasticsearch.painless.spi.Whitelist;
+import org.elasticsearch.script.ScriptContext;
+
+import java.util.List;
+import java.util.Map;
+
+public class RuntimeFieldsPainlessExtension implements PainlessExtension {
+    @Override
+    public Map<ScriptContext<?>, List<Whitelist>> getContextWhitelists() {
+        return Map.of(
+            LongScriptFieldScript.CONTEXT, LongScriptFieldScript.whitelist(),
+            StringScriptFieldScript.CONTEXT, StringScriptFieldScript.whitelist()
+        );
+    }
+}

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/StringScriptFieldScript.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+package org.elasticsearch.xpack.runtimefields;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.elasticsearch.painless.spi.Whitelist;
+import org.elasticsearch.painless.spi.WhitelistLoader;
+import org.elasticsearch.script.ScriptContext;
+import org.elasticsearch.script.ScriptFactory;
+import org.elasticsearch.search.lookup.DocLookup;
+import org.elasticsearch.search.lookup.SourceLookup;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Consumer;
+
+public abstract class StringScriptFieldScript extends AbstractScriptFieldsScript {
+    static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("string_script_field", Factory.class);
+    static List<Whitelist> whitelist() {
+        return List.of(WhitelistLoader.loadFromResourceFiles(RuntimeFieldsPainlessExtension.class, "string_whitelist.txt"));
+    }
+
+    public static final String[] PARAMETERS = {};
+
+    public interface Factory extends ScriptFactory {
+        LeafFactory newFactory(Map<String, Object> params, SourceLookup source, DocLookup fieldData);
+    }
+
+    public static interface LeafFactory {
+        StringScriptFieldScript newInstance(LeafReaderContext ctx, Consumer<String> sync) throws IOException;
+    }
+
+    private final Consumer<String> sync;
+
+    public StringScriptFieldScript(
+        Map<String, Object> params,
+        SourceLookup source,
+        DocLookup fieldData,
+        LeafReaderContext ctx,
+        Consumer<String> sync
+    ) {
+        super(params, source, fieldData, ctx);
+        this.sync = sync;
+    }
+
+    public static class Value {
+        private final StringScriptFieldScript script;
+
+        public Value(StringScriptFieldScript script) {
+            this.script = script;
+        }
+
+        public void value(String v) {
+            script.sync.accept(v);
+        }
+    }
+}

x-pack/plugin/runtime-fields/src/main/resources/META-INF/services/org.elasticsearch.painless.spi.PainlessExtension

@@ -0,0 +1 @@
+org.elasticsearch.xpack.runtimefields.RuntimeFieldsPainlessExtension

x-pack/plugin/runtime-fields/src/main/resources/org/elasticsearch/xpack/runtimefields/long_whitelist.txt

@@ -0,0 +1,14 @@
+#
+# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+# or more contributor license agreements. Licensed under the Elastic License;
+# you may not use this file except in compliance with the Elastic License.
+#
+
+# The whitelist for long-valued runtime fields
+
+class org.elasticsearch.xpack.runtimefields.LongScriptFieldScript @no_import {
+}
+
+static_import {
+    void value(org.elasticsearch.xpack.runtimefields.LongScriptFieldScript, long) bound_to org.elasticsearch.xpack.runtimefields.LongScriptFieldScript$Value
+}

x-pack/plugin/runtime-fields/src/main/resources/org/elasticsearch/xpack/runtimefields/string_whitelist.txt

@@ -0,0 +1,14 @@
+#
+# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+# or more contributor license agreements. Licensed under the Elastic License;
+# you may not use this file except in compliance with the Elastic License.
+#
+
+# The whitelist for long-valued runtime fields
+
+class org.elasticsearch.xpack.runtimefields.StringScriptFieldScript @no_import {
+}
+
+static_import {
+    void value(org.elasticsearch.xpack.runtimefields.StringScriptFieldScript, String) bound_to org.elasticsearch.xpack.runtimefields.StringScriptFieldScript$Value
+}