Skip to content

Commit c34e14c

Browse files
committed
Scripting: Expose all doc field str const accesses
* Add DocFieldsPhase visitor during compilation * Fields are accessible via ScriptScope.docFields() Doc field accesses that are expressions, even constant expressions are not handled. Refs: elastic#60001
1 parent 1781d4a commit c34e14c

File tree

5 files changed

+215
-0
lines changed

5 files changed

+215
-0
lines changed

modules/lang-painless/src/main/java/org/elasticsearch/painless/Compiler.java

+5
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import org.elasticsearch.painless.ir.ClassNode;
2525
import org.elasticsearch.painless.lookup.PainlessLookup;
2626
import org.elasticsearch.painless.node.SClass;
27+
import org.elasticsearch.painless.phase.DocFieldsPhase;
2728
import org.elasticsearch.painless.phase.SemanticHeaderPhase;
2829
import org.elasticsearch.painless.phase.UserTreeToIRTreeVisitor;
2930
import org.elasticsearch.painless.spi.Whitelist;
@@ -216,6 +217,8 @@ ScriptScope compile(Loader loader, String name, String source, CompilerSettings
216217
ScriptScope scriptScope = new ScriptScope(painlessLookup, settings, scriptClassInfo, scriptName, source, root.getIdentifier() + 1);
217218
new SemanticHeaderPhase().visitClass(root, scriptScope);
218219
root.analyze(scriptScope);
220+
// TODO(stu): Make this phase optional #60156
221+
new DocFieldsPhase().visitClass(root, scriptScope);
219222
new UserTreeToIRTreeVisitor().visitClass(root, scriptScope);
220223
ClassNode classNode = (ClassNode)scriptScope.getDecoration(root, IRNodeDecoration.class).getIRNode();
221224
DefBootstrapInjectionPhase.phase(classNode);
@@ -249,6 +252,8 @@ byte[] compile(String name, String source, CompilerSettings settings, Printer de
249252
ScriptScope scriptScope = new ScriptScope(painlessLookup, settings, scriptClassInfo, scriptName, source, root.getIdentifier() + 1);
250253
new SemanticHeaderPhase().visitClass(root, scriptScope);
251254
root.analyze(scriptScope);
255+
// TODO(stu): Make this phase optional #60156
256+
new DocFieldsPhase().visitClass(root, scriptScope);
252257
new UserTreeToIRTreeVisitor().visitClass(root, scriptScope);
253258
ClassNode classNode = (ClassNode)scriptScope.getDecoration(root, IRNodeDecoration.class).getIRNode();
254259
classNode.setDebugStream(debugStream);
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
/*
2+
* Licensed to Elasticsearch under one or more contributor
3+
* license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright
5+
* ownership. Elasticsearch licenses this file to you under
6+
* the Apache License, Version 2.0 (the "License"); you may
7+
* not use this file except in compliance with the License.
8+
* You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.elasticsearch.painless.phase;
21+
22+
import org.elasticsearch.painless.node.AExpression;
23+
import org.elasticsearch.painless.node.EBrace;
24+
import org.elasticsearch.painless.node.ECall;
25+
import org.elasticsearch.painless.node.EDot;
26+
import org.elasticsearch.painless.node.EString;
27+
import org.elasticsearch.painless.node.ESymbol;
28+
import org.elasticsearch.painless.symbol.Decorations;
29+
import org.elasticsearch.painless.symbol.ScriptScope;
30+
31+
import java.util.List;
32+
33+
/**
34+
* Find all document field accesses.
35+
*/
36+
public class DocFieldsPhase extends UserTreeBaseVisitor<ScriptScope> {
37+
@Override
38+
public void visitSymbol(ESymbol userSymbolNode, ScriptScope scriptScope) {
39+
// variables are a leaf node
40+
if (userSymbolNode.getSymbol().equals("doc")) {
41+
scriptScope.setCondition(userSymbolNode, Decorations.IsDocument.class);
42+
}
43+
}
44+
45+
@Override
46+
public void visitBrace(EBrace userBraceNode, ScriptScope scriptScope) {
47+
userBraceNode.getPrefixNode().visit(this, scriptScope);
48+
scriptScope.replicateCondition(userBraceNode.getPrefixNode(), userBraceNode.getIndexNode(), Decorations.IsDocument.class);
49+
userBraceNode.getIndexNode().visit(this, scriptScope);
50+
}
51+
52+
@Override
53+
public void visitDot(EDot userDotNode, ScriptScope scriptScope) {
54+
AExpression prefixNode = userDotNode.getPrefixNode();
55+
prefixNode.visit(this, scriptScope);
56+
if (scriptScope.getCondition(prefixNode, Decorations.IsDocument.class)) {
57+
scriptScope.addDocField(userDotNode.getIndex());
58+
}
59+
}
60+
61+
@Override
62+
public void visitCall(ECall userCallNode, ScriptScope scriptScope) {
63+
// looking for doc.get
64+
AExpression prefixNode = userCallNode.getPrefixNode();
65+
prefixNode.visit(this, scriptScope);
66+
67+
List<AExpression> argumentNodes = userCallNode.getArgumentNodes();
68+
if (argumentNodes.size() != 1 || userCallNode.getMethodName().equals("get") == false) {
69+
for (AExpression argumentNode : argumentNodes) {
70+
argumentNode.visit(this, scriptScope);
71+
}
72+
return;
73+
}
74+
75+
AExpression argument = argumentNodes.get(0);
76+
scriptScope.replicateCondition(prefixNode, argument, Decorations.IsDocument.class);
77+
argument.visit(this, scriptScope);
78+
}
79+
80+
@Override
81+
public void visitString(EString userStringNode, ScriptScope scriptScope) {
82+
if (scriptScope.getCondition(userStringNode, Decorations.IsDocument.class)) {
83+
scriptScope.addDocField(userStringNode.getString());
84+
}
85+
}
86+
}

modules/lang-painless/src/main/java/org/elasticsearch/painless/symbol/Decorations.java

+4
Original file line numberDiff line numberDiff line change
@@ -583,4 +583,8 @@ public IRNode getIRNode() {
583583
return irNode;
584584
}
585585
}
586+
587+
public interface IsDocument extends Condition {
588+
589+
}
586590
}

modules/lang-painless/src/main/java/org/elasticsearch/painless/symbol/ScriptScope.java

+14
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,10 @@
2424
import org.elasticsearch.painless.lookup.PainlessLookup;
2525
import org.elasticsearch.painless.node.ANode;
2626

27+
import java.util.ArrayList;
2728
import java.util.Collections;
2829
import java.util.HashMap;
30+
import java.util.List;
2931
import java.util.Map;
3032
import java.util.Objects;
3133
import java.util.Set;
@@ -45,6 +47,7 @@ public class ScriptScope extends Decorator {
4547
protected int syntheticCounter = 0;
4648

4749
protected boolean deterministic = true;
50+
protected List<String> docFields = new ArrayList<>();
4851
protected Set<String> usedVariables = Collections.emptySet();
4952
protected Map<String, Object> staticConstants = new HashMap<>();
5053

@@ -104,6 +107,17 @@ public boolean isDeterministic() {
104107
return deterministic;
105108
}
106109

110+
/**
111+
* Document fields read or written using constant strings
112+
*/
113+
public List<String> docFields() {
114+
return Collections.unmodifiableList(docFields);
115+
}
116+
117+
public void addDocField(String field) {
118+
docFields.add(field);
119+
}
120+
107121
public void setUsedVariables(Set<String> usedVariables) {
108122
this.usedVariables = usedVariables;
109123
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
/*
2+
* Licensed to Elasticsearch under one or more contributor
3+
* license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright
5+
* ownership. Elasticsearch licenses this file to you under
6+
* the Apache License, Version 2.0 (the "License"); you may
7+
* not use this file except in compliance with the License.
8+
* You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.elasticsearch.painless;
21+
22+
import org.elasticsearch.painless.lookup.PainlessLookup;
23+
import org.elasticsearch.painless.lookup.PainlessLookupBuilder;
24+
import org.elasticsearch.painless.spi.Whitelist;
25+
import org.elasticsearch.painless.symbol.ScriptScope;
26+
import org.elasticsearch.script.ScriptContext;
27+
28+
import java.security.AccessController;
29+
import java.security.PrivilegedAction;
30+
import java.util.Collections;
31+
import java.util.List;
32+
import java.util.Map;
33+
34+
public class DocFieldsPhaseTests extends ScriptTestCase {
35+
PainlessLookup lookup = PainlessLookupBuilder.buildFromWhitelists(Whitelist.BASE_WHITELISTS);
36+
37+
ScriptScope compile(String script) {
38+
Compiler compiler = new Compiler(
39+
MockDocTestScript.CONTEXT.instanceClazz,
40+
MockDocTestScript.CONTEXT.factoryClazz,
41+
MockDocTestScript.CONTEXT.statefulFactoryClazz, lookup
42+
);
43+
44+
// Create our loader (which loads compiled code with no permissions).
45+
final Compiler.Loader loader = AccessController.doPrivileged(new PrivilegedAction<>() {
46+
@Override
47+
public Compiler.Loader run() {
48+
return compiler.createLoader(getClass().getClassLoader());
49+
}
50+
});
51+
52+
return compiler.compile(loader,"test", script, new CompilerSettings());
53+
}
54+
55+
public abstract static class MockDocTestScript {
56+
public static final String[] PARAMETERS = {"doc", "other"};
57+
public abstract void execute(Map<String, Object> doc, Map<String, Object> other);
58+
59+
public interface Factory {
60+
MockDocTestScript newInstance();
61+
}
62+
63+
public static final ScriptContext<Factory> CONTEXT =
64+
new ScriptContext<>("test", MockDocTestScript.Factory.class);
65+
}
66+
67+
public void testArray() {
68+
List<String> expected = List.of("my_field");
69+
// Order shouldn't matter
70+
assertEquals(expected, compile("def a = doc['my_field']; def b = other['foo']").docFields());
71+
assertEquals(expected, compile("def b = other['foo']; def a = doc['my_field']").docFields());
72+
73+
// Only collect array on doc
74+
assertEquals(Collections.emptyList(), compile("def a = other['bar']").docFields());
75+
76+
// Only handle str const
77+
assertEquals(Collections.emptyList(), compile("String f = 'bar'; def a = other[f]").docFields());
78+
}
79+
80+
public void testDot() {
81+
List<String> expected = List.of("my_field");
82+
// Order shouldn't matter
83+
assertEquals(expected, compile("def a = doc.my_field; def b = other.foo").docFields());
84+
assertEquals(expected, compile("def b = other.foo; def a = doc.my_field").docFields());
85+
86+
// Only collect doc dots
87+
assertEquals(Collections.emptyList(), compile("def a = other.bar").docFields());
88+
}
89+
90+
public void testGet() {
91+
// Order shouldn't matter
92+
List<String> expected = List.of("my_field");
93+
assertEquals(expected, compile("def a = doc.get('my_field'); def b = other.get('foo')").docFields());
94+
assertEquals(expected, compile("def b = other.get('foo'); def a = doc.get('my_field')").docFields());
95+
96+
// Should work in Lambda
97+
assertEquals(expected, compile("[].sort((a, b) -> doc.get('my_field')); [].sort((a, b) -> doc.equals('bar') ? 1:2)").docFields());
98+
99+
// Only collect get on doc
100+
assertEquals(Collections.emptyList(), compile("def a = other.get('bar')").docFields());
101+
assertEquals(Collections.emptyList(), compile("def a = doc.equals('bar')").docFields());
102+
103+
// Only handle str const
104+
assertEquals(Collections.emptyList(), compile("String f = 'bar'; def b = doc.get(f)").docFields());
105+
}
106+
}

0 commit comments

Comments
 (0)