diff --git a/modules/lang-painless/src/test/java/org/elasticsearch/painless/ScriptedMetricAggContextsTests.java b/modules/lang-painless/src/test/java/org/elasticsearch/painless/ScriptedMetricAggContextsTests.java index 5c6fbc54667..2d33853b88f 100644 --- a/modules/lang-painless/src/test/java/org/elasticsearch/painless/ScriptedMetricAggContextsTests.java +++ b/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, List> 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 params = new HashMap<>(); + Map 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, diff --git a/server/src/main/java/org/elasticsearch/script/ScriptedMetricAggContexts.java b/server/src/main/java/org/elasticsearch/script/ScriptedMetricAggContexts.java index e72d597a6af..4c51b9fed69 100644 --- a/server/src/main/java/org/elasticsearch/script/ScriptedMetricAggContexts.java +++ b/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 params; private final Map state; - ParamsAndStateBase(Map params, Map state) { + public InitScript(Map params, Map 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 params, Map state) { - super(params, state); - } public abstract void execute(); @@ -64,14 +61,51 @@ public class ScriptedMetricAggContexts { public static ScriptContext CONTEXT = new ScriptContext<>("aggs_init", Factory.class); } - public abstract static class MapScript extends ParamsAndStateBase { + public abstract static class MapScript { + private static final Map DEPRECATIONS; + + static { + Map 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 params; + private final Map state; private final LeafSearchLookup leafLookup; private Scorable scorer; public MapScript(Map params, Map 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 getParams() { + return params; + } + + public Map 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 CONTEXT = new ScriptContext<>("aggs_map", Factory.class); } - public abstract static class CombineScript extends ParamsAndStateBase { + public abstract static class CombineScript { + private final Map params; + private final Map state; + public CombineScript(Map params, Map state) { - super(params, state); + this.params = params; + this.state = state; + } + + public Map getParams() { + return params; + } + + public Map getState() { + return state; } public abstract Object execute();