Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 1 commit into from Oct 24, 2018
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
Expand Up @@ -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() {
Expand Down Expand Up @@ -57,7 +66,7 @@ public void testInitBasic() {
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());

Expand All @@ -82,6 +91,32 @@ public void testMapBasic() {
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,
Expand Down
Expand Up @@ -27,15 +27,18 @@
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;
}
Expand All @@ -47,12 +50,6 @@ public Map<String, Object> getParams() {
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();

Expand All @@ -64,14 +61,51 @@ public interface Factory {
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
Expand Down Expand Up @@ -117,9 +151,21 @@ public interface Factory {
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();
Expand Down