Skip to content

Commit e398580

Browse files
jdconradjdconrad
authored
Add ability to augment classes with fields from other classes in Painless (#76628)
This change adds a a new "augmented" annotation to the Painless allowlist parser. The first use of the annotation supports adding static final fields to a specified allowlist class from another class. This supports the fields api as we can add additional fields types from other classes and augment the Field class with the new types.
1 parent 1784805 commit e398580

File tree

10 files changed

+189
-17
lines changed

10 files changed

+189
-17
lines changed

‎modules/lang-painless/spi/src/main/java/org/elasticsearch/painless/spi/annotation/AugmentedAnnotation.java‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0 and the Server Side Public License, v 1; you may not use this file except
5+
* in compliance with, at your election, the Elastic License 2.0 or the Server
6+
* Side Public License, v 1.
7+
*/
8+
9+
package org.elasticsearch.painless.spi.annotation;
10+
11+
/**
12+
* Use an augmented annotation to augment a class to have available
13+
* a static final field from a different class.
14+
*
15+
* Example:
16+
* {@code
17+
* // Whitelist
18+
*
19+
*...
20+
* class java.lang.Integer {
21+
* ...
22+
* }
23+
*
24+
* class Augmented {
25+
* int MAX_VALUE @augmented[augmented_canonical_class_name="java.lang.Integer"]
26+
* int MIN_VALUE @augmented[augmented_canonical_class_name="java.lang.Integer"]
27+
* }
28+
* ...
29+
*
30+
* // Script
31+
* ...
32+
* long value = <some_value>;
33+
* int augmented = 0;
34+
* if (value < Augmented.MAX_VALUE && value > Augmented.MIN_VALUE) {
35+
* augmented = (int)value;
36+
* }
37+
* ...
38+
* }
39+
*/
40+
public class AugmentedAnnotation {
41+
42+
public static final String NAME = "augmented";
43+
44+
private final String augmentedCanonicalClassName;
45+
46+
public AugmentedAnnotation(String augmentedCanonicalClassName) {
47+
this.augmentedCanonicalClassName = augmentedCanonicalClassName;
48+
}
49+
50+
public String getAugmentedCanonicalClassName() {
51+
return augmentedCanonicalClassName;
52+
}
53+
}

‎modules/lang-painless/spi/src/main/java/org/elasticsearch/painless/spi/annotation/AugmentedAnnotationParser.java‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0 and the Server Side Public License, v 1; you may not use this file except
5+
* in compliance with, at your election, the Elastic License 2.0 or the Server
6+
* Side Public License, v 1.
7+
*/
8+
9+
package org.elasticsearch.painless.spi.annotation;
10+
11+
import java.util.Map;
12+
13+
public class AugmentedAnnotationParser implements WhitelistAnnotationParser {
14+
15+
public static final AugmentedAnnotationParser INSTANCE = new AugmentedAnnotationParser();
16+
17+
public static final String AUGMENTED_CANONICAL_CLASS_NAME = "augmented_canonical_class_name";
18+
19+
private AugmentedAnnotationParser() {
20+
21+
}
22+
23+
@Override
24+
public Object parse(Map<String, String> arguments) {
25+
String javaClassName = arguments.get(AUGMENTED_CANONICAL_CLASS_NAME);
26+
27+
if (javaClassName == null) {
28+
throw new IllegalArgumentException("augmented_canonical_class_name cannot be null for [@augmented] annotation");
29+
}
30+
31+
if ((arguments.isEmpty() || arguments.size() == 1 && arguments.containsKey(AUGMENTED_CANONICAL_CLASS_NAME)) == false) {
32+
throw new IllegalArgumentException("unexpected parameters for [@augmented] annotation, found " + arguments);
33+
}
34+
35+
return new AugmentedAnnotation(javaClassName);
36+
}
37+
}

‎modules/lang-painless/spi/src/main/java/org/elasticsearch/painless/spi/annotation/WhitelistAnnotationParser.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,8 @@ public interface WhitelistAnnotationParser {
2626
new AbstractMap.SimpleEntry<>(DeprecatedAnnotation.NAME, DeprecatedAnnotationParser.INSTANCE),
2727
new AbstractMap.SimpleEntry<>(NonDeterministicAnnotation.NAME, NonDeterministicAnnotationParser.INSTANCE),
2828
new AbstractMap.SimpleEntry<>(InjectConstantAnnotation.NAME, InjectConstantAnnotationParser.INSTANCE),
29-
new AbstractMap.SimpleEntry<>(CompileTimeOnlyAnnotation.NAME, CompileTimeOnlyAnnotationParser.INSTANCE)
29+
new AbstractMap.SimpleEntry<>(CompileTimeOnlyAnnotation.NAME, CompileTimeOnlyAnnotationParser.INSTANCE),
30+
new AbstractMap.SimpleEntry<>(AugmentedAnnotation.NAME, AugmentedAnnotationParser.INSTANCE)
3031
).collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))
3132
);
3233

‎modules/lang-painless/spi/src/test/java/org/elasticsearch/painless/WhitelistLoaderTests.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import java.util.Map;
2222

2323
public class WhitelistLoaderTests extends ESTestCase {
24+
2425
public void testUnknownAnnotations() {
2526
Map<String, WhitelistAnnotationParser> parsers = new HashMap<>(WhitelistAnnotationParser.BASE_ANNOTATION_PARSERS);
2627

‎modules/lang-painless/src/main/java/org/elasticsearch/painless/lookup/PainlessField.java‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,19 +10,24 @@
1010

1111
import java.lang.invoke.MethodHandle;
1212
import java.lang.reflect.Field;
13+
import java.util.Map;
1314
import java.util.Objects;
1415

1516
public final class PainlessField {
1617

1718
public final Field javaField;
1819
public final Class<?> typeParameter;
20+
public final Map<Class<?>, Object> annotations;
1921

2022
public final MethodHandle getterMethodHandle;
2123
public final MethodHandle setterMethodHandle;
2224

23-
PainlessField(Field javaField, Class<?> typeParameter, MethodHandle getterMethodHandle, MethodHandle setterMethodHandle) {
25+
PainlessField(Field javaField, Class<?> typeParameter, Map<Class<?>, Object> annotations,
26+
MethodHandle getterMethodHandle, MethodHandle setterMethodHandle) {
27+
2428
this.javaField = javaField;
2529
this.typeParameter = typeParameter;
30+
this.annotations = annotations;
2631

2732
this.getterMethodHandle = getterMethodHandle;
2833
this.setterMethodHandle = setterMethodHandle;
@@ -41,11 +46,12 @@ public boolean equals(Object object) {
4146
PainlessField that = (PainlessField)object;
4247

4348
return Objects.equals(javaField, that.javaField) &&
44-
Objects.equals(typeParameter, that.typeParameter);
49+
Objects.equals(typeParameter, that.typeParameter) &&
50+
Objects.equals(annotations, that.annotations);
4551
}
4652

4753
@Override
4854
public int hashCode() {
49-
return Objects.hash(javaField, typeParameter);
55+
return Objects.hash(javaField, typeParameter, annotations);
5056
}
5157
}

‎modules/lang-painless/src/main/java/org/elasticsearch/painless/lookup/PainlessLookupBuilder.java‎

Lines changed: 52 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import org.elasticsearch.painless.spi.WhitelistField;
2020
import org.elasticsearch.painless.spi.WhitelistInstanceBinding;
2121
import org.elasticsearch.painless.spi.WhitelistMethod;
22+
import org.elasticsearch.painless.spi.annotation.AugmentedAnnotation;
2223
import org.elasticsearch.painless.spi.annotation.CompileTimeOnlyAnnotation;
2324
import org.elasticsearch.painless.spi.annotation.InjectConstantAnnotation;
2425
import org.elasticsearch.painless.spi.annotation.NoImportAnnotation;
@@ -140,7 +141,8 @@ public static PainlessLookup buildFromWhitelists(List<Whitelist> whitelists) {
140141
for (WhitelistField whitelistField : whitelistClass.whitelistFields) {
141142
origin = whitelistField.origin;
142143
painlessLookupBuilder.addPainlessField(
143-
targetCanonicalClassName, whitelistField.fieldName, whitelistField.canonicalTypeNameParameter);
144+
whitelist.classLoader, targetCanonicalClassName,
145+
whitelistField.fieldName, whitelistField.canonicalTypeNameParameter, whitelistField.painlessAnnotations);
144146
}
145147
}
146148

@@ -426,6 +428,7 @@ public void addPainlessMethod(ClassLoader classLoader, String targetCanonicalCla
426428
Objects.requireNonNull(methodName);
427429
Objects.requireNonNull(returnCanonicalTypeName);
428430
Objects.requireNonNull(canonicalTypeNameParameters);
431+
Objects.requireNonNull(annotations);
429432

430433
Class<?> targetClass = canonicalClassNamesToClasses.get(targetCanonicalClassName);
431434

@@ -472,6 +475,7 @@ public void addPainlessMethod(Class<?> targetClass, Class<?> augmentedClass,
472475
Objects.requireNonNull(methodName);
473476
Objects.requireNonNull(returnType);
474477
Objects.requireNonNull(typeParameters);
478+
Objects.requireNonNull(annotations);
475479

476480
if (targetClass == def.class) {
477481
throw new IllegalArgumentException("cannot add method to reserved class [" + DEF_CLASS_NAME + "]");
@@ -610,10 +614,14 @@ public void addPainlessMethod(Class<?> targetClass, Class<?> augmentedClass,
610614
}
611615
}
612616

613-
public void addPainlessField(String targetCanonicalClassName, String fieldName, String canonicalTypeNameParameter) {
617+
public void addPainlessField(ClassLoader classLoader, String targetCanonicalClassName,
618+
String fieldName, String canonicalTypeNameParameter, Map<Class<?>, Object> annotations) {
619+
620+
Objects.requireNonNull(classLoader);
614621
Objects.requireNonNull(targetCanonicalClassName);
615622
Objects.requireNonNull(fieldName);
616623
Objects.requireNonNull(canonicalTypeNameParameter);
624+
Objects.requireNonNull(annotations);
617625

618626
Class<?> targetClass = canonicalClassNamesToClasses.get(targetCanonicalClassName);
619627

@@ -622,6 +630,17 @@ public void addPainlessField(String targetCanonicalClassName, String fieldName,
622630
"[[" + targetCanonicalClassName + "], [" + fieldName + "], [" + canonicalTypeNameParameter + "]]");
623631
}
624632

633+
String augmentedCanonicalClassName = annotations.containsKey(AugmentedAnnotation.class) ?
634+
((AugmentedAnnotation)annotations.get(AugmentedAnnotation.class)).getAugmentedCanonicalClassName() : null;
635+
636+
Class<?> augmentedClass = null;
637+
638+
if (augmentedCanonicalClassName != null) {
639+
augmentedClass = loadClass(classLoader, augmentedCanonicalClassName,
640+
() -> "augmented class [" + augmentedCanonicalClassName + "] not found for field " +
641+
"[[" + targetCanonicalClassName + "], [" + fieldName + "]");
642+
}
643+
625644
Class<?> typeParameter = canonicalTypeNameToType(canonicalTypeNameParameter);
626645

627646
if (typeParameter == null) {
@@ -630,13 +649,16 @@ public void addPainlessField(String targetCanonicalClassName, String fieldName,
630649
}
631650

632651

633-
addPainlessField(targetClass, fieldName, typeParameter);
652+
addPainlessField(targetClass, augmentedClass, fieldName, typeParameter, annotations);
634653
}
635654

636-
public void addPainlessField(Class<?> targetClass, String fieldName, Class<?> typeParameter) {
655+
public void addPainlessField(Class<?> targetClass, Class<?> augmentedClass,
656+
String fieldName, Class<?> typeParameter, Map<Class<?>, Object> annotations) {
657+
637658
Objects.requireNonNull(targetClass);
638659
Objects.requireNonNull(fieldName);
639660
Objects.requireNonNull(typeParameter);
661+
Objects.requireNonNull(annotations);
640662

641663
if (targetClass == def.class) {
642664
throw new IllegalArgumentException("cannot add field to reserved class [" + DEF_CLASS_NAME + "]");
@@ -658,17 +680,33 @@ public void addPainlessField(Class<?> targetClass, String fieldName, Class<?> ty
658680
}
659681

660682
if (isValidType(typeParameter) == false) {
661-
throw new IllegalArgumentException("type parameter [" + typeToCanonicalTypeName(typeParameter) + "] not found " +
662-
"for field [[" + targetCanonicalClassName + "], [" + fieldName + "]");
683+
throw new IllegalArgumentException("type parameter [" + typeToCanonicalTypeName(typeParameter) + "] not found for field " +
684+
"[[" + targetCanonicalClassName + "], [" + fieldName + "], [" + typeToCanonicalTypeName(typeParameter) + "]]");
663685
}
664686

665687
Field javaField;
666688

667-
try {
668-
javaField = targetClass.getField(fieldName);
669-
} catch (NoSuchFieldException nsme) {
670-
throw new IllegalArgumentException(
671-
"reflection object not found for field [[" + targetCanonicalClassName + "], [" + fieldName + "]", nsme);
689+
if (augmentedClass == null) {
690+
try {
691+
javaField = targetClass.getField(fieldName);
692+
} catch (NoSuchFieldException nsfe) {
693+
throw new IllegalArgumentException("reflection object not found for field " +
694+
"[[" + targetCanonicalClassName + "], [" + fieldName + "], [" + typeToCanonicalTypeName(typeParameter) + "]]",
695+
nsfe);
696+
}
697+
} else {
698+
try {
699+
javaField = augmentedClass.getField(fieldName);
700+
701+
if (Modifier.isStatic(javaField.getModifiers()) == false || Modifier.isFinal(javaField.getModifiers()) == false) {
702+
throw new IllegalArgumentException("field [[" + targetCanonicalClassName + "], [" + fieldName + "] " +
703+
"with augmented class [" + typeToCanonicalTypeName(augmentedClass) + "] must be static and final");
704+
}
705+
} catch (NoSuchFieldException nsfe) {
706+
throw new IllegalArgumentException("reflection object not found for field " +
707+
"[[" + targetCanonicalClassName + "], [" + fieldName + "], [" + typeToCanonicalTypeName(typeParameter) + "]]" +
708+
"with augmented class [" + typeToCanonicalTypeName(augmentedClass) + "]", nsfe);
709+
}
672710
}
673711

674712
if (javaField.getType() != typeToJavaType(typeParameter)) {
@@ -694,7 +732,7 @@ public void addPainlessField(Class<?> targetClass, String fieldName, Class<?> ty
694732
}
695733

696734
PainlessField existingPainlessField = painlessClassBuilder.staticFields.get(painlessFieldKey);
697-
PainlessField newPainlessField = new PainlessField(javaField, typeParameter, methodHandleGetter, null);
735+
PainlessField newPainlessField = new PainlessField(javaField, typeParameter, annotations, methodHandleGetter, null);
698736

699737
if (existingPainlessField == null) {
700738
newPainlessField = painlessFieldCache.computeIfAbsent(newPainlessField, key -> key);
@@ -718,7 +756,8 @@ public void addPainlessField(Class<?> targetClass, String fieldName, Class<?> ty
718756
}
719757

720758
PainlessField existingPainlessField = painlessClassBuilder.fields.get(painlessFieldKey);
721-
PainlessField newPainlessField = new PainlessField(javaField, typeParameter, methodHandleGetter, methodHandleSetter);
759+
PainlessField newPainlessField =
760+
new PainlessField(javaField, typeParameter, annotations, methodHandleGetter, methodHandleSetter);
722761

723762
if (existingPainlessField == null) {
724763
newPainlessField = painlessFieldCache.computeIfAbsent(newPainlessField, key -> key);

‎modules/lang-painless/src/test/java/org/elasticsearch/painless/AugmentationTests.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -281,4 +281,11 @@ public void testToEpochMilli() {
281281
assertEquals(0L, exec("ZonedDateTime.parse('1970-01-01T00:00:00Z').toEpochMilli()"));
282282
assertEquals(1602097376782L, exec("ZonedDateTime.parse('2020-10-07T19:02:56.782Z').toEpochMilli()"));
283283
}
284+
285+
public void testAugmentedField() {
286+
assertEquals(Integer.MAX_VALUE, exec("org.elasticsearch.painless.FeatureTestObject.MAX_VALUE"));
287+
assertEquals("test_string", exec("Integer.STRINGS[0]"));
288+
assertEquals("test_value", exec("Map.STRINGS['test_key']"));
289+
assertTrue((boolean)exec("Integer.STRINGS[0].substring(0, 4) == Map.STRINGS['test_key'].substring(0, 4)"));
290+
}
284291
}

‎modules/lang-painless/src/test/java/org/elasticsearch/painless/FeatureTestObject.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package org.elasticsearch.painless;
22

3+
import java.util.Collections;
34
import java.util.List;
45
import java.util.function.Function;
56

@@ -37,6 +38,8 @@ public static int staticNumberArgument(int injected, int userArgument) {
3738
return injected * userArgument;
3839
}
3940

41+
public static final List<String> STRINGS = Collections.singletonList("test_string");
42+
4043
private int x;
4144
private int y;
4245
public int z;

‎modules/lang-painless/src/test/java/org/elasticsearch/painless/FeatureTestObject3.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0 and the Server Side Public License, v 1; you may not use this file except
5+
* in compliance with, at your election, the Elastic License 2.0 or the Server
6+
* Side Public License, v 1.
7+
*/
8+
9+
package org.elasticsearch.painless;
10+
11+
import java.util.Collections;
12+
import java.util.Map;
13+
14+
public class FeatureTestObject3 {
15+
16+
public static final Map<String, String> STRINGS = Collections.singletonMap("test_key", "test_value");
17+
}

‎modules/lang-painless/src/test/resources/org/elasticsearch/painless/org.elasticsearch.painless.test‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,16 @@ class org.elasticsearch.painless.TestFieldScript @no_import {
2121
class org.elasticsearch.painless.TestFieldScript$Factory @no_import {
2222
}
2323

24+
class java.lang.Integer {
25+
List STRINGS @augmented[augmented_canonical_class_name="org.elasticsearch.painless.FeatureTestObject"]
26+
}
27+
28+
class java.util.Map {
29+
Map STRINGS @augmented[augmented_canonical_class_name="org.elasticsearch.painless.FeatureTestObject3"]
30+
}
2431

2532
class org.elasticsearch.painless.FeatureTestObject @no_import {
33+
int MAX_VALUE @augmented[augmented_canonical_class_name="java.lang.Integer"]
2634
int z
2735
()
2836
(int,int)

0 commit comments

Comments
 (0)