Scripting: Add back params._source access in scripted metric aggs (#34777)

Access to special variables _source and _fields were accidentally
removed in recent refactorings. This commit adds them back, along with a
test.

closes #33884

modules/lang-painless/src/test/java/org/elasticsearch/painless/ScriptedMetricAggContextsTests.java

@@ -19,17 +19,26 @@
 
 package org.elasticsearch.painless;
 
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.memory.MemoryIndex;
 import org.apache.lucene.search.Scorable;
 import org.elasticsearch.painless.spi.Whitelist;
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.ScriptedMetricAggContexts;
+import org.elasticsearch.search.lookup.LeafSearchLookup;
+import org.elasticsearch.search.lookup.SearchLookup;
+import org.elasticsearch.search.lookup.SourceLookup;
 
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
 public class ScriptedMetricAggContextsTests extends ScriptTestCase {
     @Override
     protected Map<ScriptContext<?>, List<Whitelist>> scriptContexts() {
@@ -57,7 +66,7 @@ public class ScriptedMetricAggContextsTests extends ScriptTestCase {
         assertEquals(10, state.get("testField"));
     }
 
-    public void testMapBasic() {
+    public void testMapBasic() throws IOException {
         ScriptedMetricAggContexts.MapScript.Factory factory = scriptEngine.compile("test",
             "state.testField = 2*_score", ScriptedMetricAggContexts.MapScript.CONTEXT, Collections.emptyMap());
 
@@ -82,6 +91,32 @@ public class ScriptedMetricAggContextsTests extends ScriptTestCase {
         assertEquals(1.0, state.get("testField"));
     }
 
+    public void testMapSourceAccess() throws IOException {
+        ScriptedMetricAggContexts.MapScript.Factory factory = scriptEngine.compile("test",
+            "state.testField = params._source.three", ScriptedMetricAggContexts.MapScript.CONTEXT, Collections.emptyMap());
+
+        Map<String, Object> params = new HashMap<>();
+        Map<String, Object> state = new HashMap<>();
+
+        MemoryIndex index = new MemoryIndex();
+        // we don't need a real index, just need to construct a LeafReaderContext which cannot be mocked
+        LeafReaderContext leafReaderContext = index.createSearcher().getIndexReader().leaves().get(0);
+
+        SearchLookup lookup = mock(SearchLookup.class);
+        LeafSearchLookup leafLookup = mock(LeafSearchLookup.class);
+        when(lookup.getLeafSearchLookup(leafReaderContext)).thenReturn(leafLookup);
+        SourceLookup sourceLookup = mock(SourceLookup.class);
+        when(leafLookup.asMap()).thenReturn(Collections.singletonMap("_source", sourceLookup));
+        when(sourceLookup.get("three")).thenReturn(3);
+        ScriptedMetricAggContexts.MapScript.LeafFactory leafFactory = factory.newFactory(params, state, lookup);
+        ScriptedMetricAggContexts.MapScript script = leafFactory.newInstance(leafReaderContext);
+
+        script.execute();
+
+        assert(state.containsKey("testField"));
+        assertEquals(3, state.get("testField"));
+    }
+
     public void testCombineBasic() {
         ScriptedMetricAggContexts.CombineScript.Factory factory = scriptEngine.compile("test",
             "state.testField = params.initialVal; return state.testField + params.inc", ScriptedMetricAggContexts.CombineScript.CONTEXT,

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

@@ -27,15 +27,18 @@ import org.elasticsearch.search.lookup.LeafSearchLookup;
 import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
 public class ScriptedMetricAggContexts {
-    private abstract static class ParamsAndStateBase {
+
+    public abstract static class InitScript {
         private final Map<String, Object> params;
         private final Map<String, Object> state;
 
-        ParamsAndStateBase(Map<String, Object> params, Map<String, Object> state) {
+        public InitScript(Map<String, Object> params, Map<String, Object> state) {
             this.params = params;
             this.state = state;
         }
@@ -47,12 +50,6 @@ public class ScriptedMetricAggContexts {
         public Object getState() {
             return state;
         }
-    }
-
-    public abstract static class InitScript extends ParamsAndStateBase {
-        public InitScript(Map<String, Object> params, Map<String, Object> state) {
-            super(params, state);
-        }
 
         public abstract void execute();
 
@@ -64,14 +61,51 @@ public class ScriptedMetricAggContexts {
         public static ScriptContext<Factory> CONTEXT = new ScriptContext<>("aggs_init", Factory.class);
     }
 
-    public abstract static class MapScript extends ParamsAndStateBase {
+    public abstract static class MapScript {
+        private static final Map<String, String> DEPRECATIONS;
+
+        static {
+            Map<String, String> deprecations = new HashMap<>();
+            deprecations.put(
+                "doc",
+                "Accessing variable [doc] via [params.doc] from within a scripted metric agg map script " +
+                    "is deprecated in favor of directly accessing [doc]."
+            );
+            deprecations.put(
+                "_doc",
+                "Accessing variable [doc] via [params._doc] from within a scripted metric agg map script " +
+                    "is deprecated in favor of directly accessing [doc]."
+            );
+            deprecations.put(
+                "_agg",
+                "Accessing variable [_agg] via [params._agg] from within a scripted metric agg map script " +
+                    "is deprecated in favor of using [state]."
+            );
+            DEPRECATIONS = Collections.unmodifiableMap(deprecations);
+        }
+
+        private final Map<String, Object> params;
+        private final Map<String, Object> state;
         private final LeafSearchLookup leafLookup;
         private Scorable scorer;
 
         public MapScript(Map<String, Object> params, Map<String, Object> state, SearchLookup lookup, LeafReaderContext leafContext) {
-            super(params, state);
+            this.state = state;
-
             this.leafLookup = leafContext == null ? null : lookup.getLeafSearchLookup(leafContext);
+            if (leafLookup != null) {
+                params = new HashMap<>(params); // copy params so we aren't modifying input
+                params.putAll(leafLookup.asMap()); // add lookup vars
+                params = new ParameterMap(params, DEPRECATIONS); // wrap with deprecations
+            }
+            this.params = params;
+        }
+
+        public Map<String, Object> getParams() {
+            return params;
+        }
+
+        public Map<String, Object> getState() {
+            return state;
         }
 
         // Return the doc as a map (instead of LeafDocLookup) in order to abide by type whitelisting rules for
@@ -117,9 +151,21 @@ public class ScriptedMetricAggContexts {
         public static ScriptContext<Factory> CONTEXT = new ScriptContext<>("aggs_map", Factory.class);
     }
 
-    public abstract static class CombineScript extends ParamsAndStateBase {
+    public abstract static class CombineScript {
+        private final Map<String, Object> params;
+        private final Map<String, Object> state;
+
         public CombineScript(Map<String, Object> params, Map<String, Object> state) {
-            super(params, state);
+            this.params = params;
+            this.state = state;
+        }
+
+        public Map<String, Object> getParams() {
+            return params;
+        }
+
+        public Map<String, Object> getState() {
+            return state;
         }
 
         public abstract Object execute();