Scripting: Replace Update Context

modules/lang-painless/src/test/resources/rest-api-spec/test/painless/15_update.yml

@@ -132,7 +132,7 @@
           body:
             script:
               lang:   painless
-              source: "for (def key : params.keySet()) { ctx._source[key] = params[key]}"
+              source: "ctx._source.ctx = ctx"
               params: { bar: 'xxx' }
 
   - match: { error.root_cause.0.type: "remote_transport_exception" }

modules/reindex/src/main/java/org/elasticsearch/index/reindex/AbstractAsyncBulkByScrollAction.java

@@ -48,9 +48,9 @@
 import org.elasticsearch.index.mapper.TypeFieldMapper;
 import org.elasticsearch.index.mapper.VersionFieldMapper;
 import org.elasticsearch.index.reindex.ScrollableHitSource.SearchFailure;
-import org.elasticsearch.script.ExecutableScript;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptService;
+import org.elasticsearch.script.UpdateScript;
 import org.elasticsearch.search.sort.SortBuilder;
 import org.elasticsearch.threadpool.ThreadPool;
 
@@ -746,7 +746,7 @@ public abstract static class ScriptApplier implements BiFunction<RequestWrapper<
         private final Script script;
         private final Map<String, Object> params;
 
-        private ExecutableScript executable;
+        private UpdateScript executable;
         private Map<String, Object> context;
 
         public ScriptApplier(WorkerBulkByScrollTaskState taskWorker,
@@ -766,7 +766,7 @@ public RequestWrapper<?> apply(RequestWrapper<?> request, ScrollableHitSource.Hi
                 return request;
             }
             if (executable == null) {
-                ExecutableScript.Factory factory = scriptService.compile(script, ExecutableScript.UPDATE_CONTEXT);
+                UpdateScript.Factory factory = scriptService.compile(script, UpdateScript.CONTEXT);
                 executable = factory.newInstance(params);
             }
             if (context == null) {
@@ -787,8 +787,7 @@ public RequestWrapper<?> apply(RequestWrapper<?> request, ScrollableHitSource.Hi
             OpType oldOpType = OpType.INDEX;
             context.put("op", oldOpType.toString());
 
-            executable.setNextVar("ctx", context);
-            executable.run();
+            executable.execute(context);
 
             String newOp = (String) context.remove("op");
             if (newOp == null) {

modules/reindex/src/test/java/org/elasticsearch/index/reindex/AbstractAsyncBulkByScrollActionScriptTestCase.java

@@ -26,8 +26,10 @@
 import org.elasticsearch.action.index.IndexRequest;
 import org.elasticsearch.script.ExecutableScript;
 import org.elasticsearch.script.ScriptService;
+import org.elasticsearch.script.UpdateScript;
 import org.junit.Before;
 
+import java.util.Collections;
 import java.util.Map;
 import java.util.function.Consumer;
 
@@ -54,10 +56,16 @@ public void setupScriptService() {
     protected <T extends ActionRequest> T applyScript(Consumer<Map<String, Object>> scriptBody) {
         IndexRequest index = new IndexRequest("index", "type", "1").source(singletonMap("foo", "bar"));
         ScrollableHitSource.Hit doc = new ScrollableHitSource.BasicHit("test", "type", "id", 0);
-        ExecutableScript executableScript = new SimpleExecutableScript(scriptBody);
-        ExecutableScript.Factory factory = params -> executableScript;
-        when(scriptService.compile(any(), eq(ExecutableScript.CONTEXT))).thenReturn(factory);
-        when(scriptService.compile(any(), eq(ExecutableScript.UPDATE_CONTEXT))).thenReturn(factory);
+        UpdateScript executableScript = new UpdateScript(Collections.emptyMap()) {
+            @Override
+            public void execute(Map<String, Object> ctx) {
+                scriptBody.accept(ctx);
+            }
+        };
+        UpdateScript.Factory factory = params -> executableScript;
+        ExecutableScript simpleExecutableScript = new SimpleExecutableScript(scriptBody);
+        when(scriptService.compile(any(), eq(ExecutableScript.CONTEXT))).thenReturn(params -> simpleExecutableScript);
+        when(scriptService.compile(any(), eq(UpdateScript.CONTEXT))).thenReturn(factory);
         AbstractAsyncBulkByScrollAction<Request> action = action(scriptService, request().setScript(mockScript("")));
         RequestWrapper<?> result = action.buildScriptApplier().apply(AbstractAsyncBulkByScrollAction.wrap(index), doc);
         return (result != null) ? (T) result.self() : null;

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

@@ -42,9 +42,9 @@
 import org.elasticsearch.index.mapper.RoutingFieldMapper;
 import org.elasticsearch.index.shard.IndexShard;
 import org.elasticsearch.index.shard.ShardId;
-import org.elasticsearch.script.ExecutableScript;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptService;
+import org.elasticsearch.script.UpdateScript;
 import org.elasticsearch.search.lookup.SourceLookup;
 
 import java.io.IOException;
@@ -279,10 +279,9 @@ Result prepareUpdateScriptRequest(ShardId shardId, UpdateRequest request, GetRes
     private Map<String, Object> executeScript(Script script, Map<String, Object> ctx) {
         try {
             if (scriptService != null) {
-                ExecutableScript.Factory factory = scriptService.compile(script, ExecutableScript.UPDATE_CONTEXT);
-                ExecutableScript executableScript = factory.newInstance(script.getParams());
-                executableScript.setNextVar(ContextFields.CTX, ctx);
-                executableScript.run();
+                UpdateScript.Factory factory = scriptService.compile(script, UpdateScript.CONTEXT);
+                UpdateScript executableScript = factory.newInstance(script.getParams());
+                executableScript.execute(ctx);
             }
         } catch (Exception e) {
             throw new IllegalArgumentException("failed to execute script", e);

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

@@ -49,5 +49,4 @@ interface Factory {
 
     // TODO: remove these once each has its own script interface
     ScriptContext<Factory> AGGS_CONTEXT = new ScriptContext<>("aggs_executable", Factory.class);
-    ScriptContext<Factory> UPDATE_CONTEXT = new ScriptContext<>("update", Factory.class);
 }

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

@@ -50,7 +50,7 @@ public class ScriptModule {
             SearchScript.TERMS_SET_QUERY_CONTEXT,
             ExecutableScript.CONTEXT,
             ExecutableScript.AGGS_CONTEXT,
-            ExecutableScript.UPDATE_CONTEXT,
+            UpdateScript.CONTEXT,
             IngestScript.CONTEXT,
             FilterScript.CONTEXT,
             SimilarityScript.CONTEXT,

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

@@ -285,7 +285,7 @@ public <FactoryType> FactoryType compile(Script script, ScriptContext<FactoryTyp
         // 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(lang);
-        boolean notSupported = context.name.equals(ExecutableScript.UPDATE_CONTEXT.name);
+        boolean notSupported = context.name.equals(UpdateScript.CONTEXT.name);
         if (expression && notSupported) {
             throw new UnsupportedOperationException("scripts of type [" + script.getType() + "]," +
                 " operation [" + context.name + "] and lang [" + lang + "] are not supported");

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

@@ -0,0 +1,52 @@
+
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.script;
+
+import java.util.Map;
+
+/**
+ * An update script.
+ */
+public abstract class UpdateScript {
+
+    public static final String[] PARAMETERS = { "ctx" };
+
+    /** The context used to compile {@link UpdateScript} factories. */
+    public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("update", Factory.class);
+
+    /** The generic runtime parameters for the script. */
+    private final Map<String, Object> params;
+
+    public UpdateScript(Map<String, Object> params) {
+        this.params = params;
+    }
+
+    /** Return the parameters for this script. */
+    public Map<String, Object> getParams() {
+        return params;
+    }
+
+    public abstract void execute(Map<String, Object> ctx);
+
+    public interface Factory {
+        UpdateScript newInstance(Map<String, Object> params);
+    }
+}

server/src/test/java/org/elasticsearch/script/ScriptServiceTests.java

@@ -167,7 +167,7 @@ public void testAllowAllScriptContextSettings() throws IOException {
 
         assertCompileAccepted("painless", "script", ScriptType.INLINE, SearchScript.CONTEXT);
         assertCompileAccepted("painless", "script", ScriptType.INLINE, SearchScript.AGGS_CONTEXT);
-        assertCompileAccepted("painless", "script", ScriptType.INLINE, ExecutableScript.UPDATE_CONTEXT);
+        assertCompileAccepted("painless", "script", ScriptType.INLINE, UpdateScript.CONTEXT);
         assertCompileAccepted("painless", "script", ScriptType.INLINE, IngestScript.CONTEXT);
     }
 
@@ -187,7 +187,7 @@ public void testAllowSomeScriptContextSettings() throws IOException {
 
         assertCompileAccepted("painless", "script", ScriptType.INLINE, SearchScript.CONTEXT);
         assertCompileAccepted("painless", "script", ScriptType.INLINE, SearchScript.AGGS_CONTEXT);
-        assertCompileRejected("painless", "script", ScriptType.INLINE, ExecutableScript.UPDATE_CONTEXT);
+        assertCompileRejected("painless", "script", ScriptType.INLINE, UpdateScript.CONTEXT);
     }
 
     public void testAllowNoScriptTypeSettings() throws IOException {

test/framework/src/main/java/org/elasticsearch/script/MockScriptEngine.java

@@ -96,6 +96,18 @@ public void execute(Map<String, Object> ctx) {
                 }
             };
             return context.factoryClazz.cast(factory);
+        } else if (context.instanceClazz.equals(UpdateScript.class)) {
+            UpdateScript.Factory factory = parameters -> new UpdateScript(parameters) {
+                @Override
+                public void execute(Map<String, Object> ctx) {
+                    final Map<String, Object> vars = new HashMap<>();
+                    vars.put("ctx", ctx);
+                    vars.put("params", parameters);
+                    vars.putAll(parameters);
+                    script.apply(vars);
+                }
+            };
+            return context.factoryClazz.cast(factory);
         } else if (context.instanceClazz.equals(TemplateScript.class)) {
             TemplateScript.Factory factory = vars -> {
                 // TODO: need a better way to implement all these new contexts