Skip to content

Commit e16bb4a

Browse files
craigtavernercraigtaverner
authored
Output function signature license requirements to Kibana definitions (#127717) (#127802)
Output function signature license requirements to Kibana definition files, and also test that this matches the actual licensing behaviour of the functions. ES|QL functions that enforce license checks do so with the `LicenseAware` interface. This does not expose what that functions license level is, but only whether the current active license will be sufficient for that function and its current signature (data types passed in as fields). Rather than add to this interface, we've made the license level information test-only information. This means if a function implements LicenseAware, it also needs to add a method to its test class to specify the license level for the signature being called. All functions will be tested for compliance, so failing to add this will result in test failure. Also if the test license level does not match the enforced license, that will also cause a failure.
1 parent 69d5b9a commit e16bb4a

File tree

15 files changed

+227
-29
lines changed

15 files changed

+227
-29
lines changed

‎docs/reference/query-languages/esql/kibana/definition/functions/categorize.json‎

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎docs/reference/query-languages/esql/kibana/definition/functions/st_extent_agg.json‎

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/DataType.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -568,6 +568,10 @@ public static boolean isSpatialPoint(DataType t) {
568568
return t == GEO_POINT || t == CARTESIAN_POINT;
569569
}
570570

571+
public static boolean isSpatialShape(DataType t) {
572+
return t == GEO_SHAPE || t == CARTESIAN_SHAPE;
573+
}
574+
571575
public static boolean isSpatialGeo(DataType t) {
572576
return t == GEO_POINT || t == GEO_SHAPE;
573577
}

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java‎

Lines changed: 103 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,9 +24,13 @@
2424
import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
2525
import org.elasticsearch.compute.test.TestBlockFactory;
2626
import org.elasticsearch.indices.CrankyCircuitBreakerService;
27+
import org.elasticsearch.license.License;
28+
import org.elasticsearch.license.XPackLicenseState;
29+
import org.elasticsearch.license.internal.XPackLicenseStatus;
2730
import org.elasticsearch.logging.LogManager;
2831
import org.elasticsearch.logging.Logger;
2932
import org.elasticsearch.test.ESTestCase;
33+
import org.elasticsearch.xpack.esql.LicenseAware;
3034
import org.elasticsearch.xpack.esql.core.expression.Attribute;
3135
import org.elasticsearch.xpack.esql.core.expression.Expression;
3236
import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
@@ -51,7 +55,6 @@
5155
import org.junit.After;
5256
import org.junit.AfterClass;
5357

54-
import java.io.IOException;
5558
import java.lang.reflect.Constructor;
5659
import java.lang.reflect.InvocationTargetException;
5760
import java.lang.reflect.Method;
@@ -701,7 +704,8 @@ public void testSerializationOfSimple() {
701704
*/
702705
@AfterClass
703706
public static void testFunctionInfo() {
704-
Logger log = LogManager.getLogger(getTestClass());
707+
Class<?> testClass = getTestClass();
708+
Logger log = LogManager.getLogger(testClass);
705709
FunctionDefinition definition = definition(functionName());
706710
if (definition == null) {
707711
log.info("Skipping function info checks because the function isn't registered");
@@ -724,7 +728,7 @@ public static void testFunctionInfo() {
724728
for (int i = 0; i < args.size(); i++) {
725729
typesFromSignature.add(new HashSet<>());
726730
}
727-
for (Map.Entry<List<DataType>, DataType> entry : signatures(getTestClass()).entrySet()) {
731+
for (Map.Entry<List<DataType>, DataType> entry : signatures(testClass).entrySet()) {
728732
List<DataType> types = entry.getKey();
729733
for (int i = 0; i < args.size() && i < types.size(); i++) {
730734
typesFromSignature.get(i).add(types.get(i).esNameIfPossible());
@@ -767,6 +771,101 @@ public static void testFunctionInfo() {
767771
assertEquals(returnFromSignature, returnTypes);
768772
}
769773

774+
/**
775+
* This test is meant to validate that the license checks documented match those enforced.
776+
* The expectations are set in the test class using a method with this signature:
777+
* <code>
778+
* public static License.OperationMode licenseRequirement(List&lt;DataType&gt; fieldTypes);
779+
* </code>
780+
* License enforcement in the function class is achieved using the interface <code>LicenseAware</code>.
781+
* This test will make sure the two are in agreement, and does not require that the function class actually
782+
* report its license level. If we add license checks to any function, but fail to also add the expected
783+
* license level to the test class, this test will fail.
784+
*/
785+
@AfterClass
786+
public static void testFunctionLicenseChecks() throws Exception {
787+
Class<?> testClass = getTestClass();
788+
Logger log = LogManager.getLogger(testClass);
789+
FunctionDefinition definition = definition(functionName());
790+
if (definition == null) {
791+
log.info("Skipping function info checks because the function isn't registered");
792+
return;
793+
}
794+
log.info("Running function license checks");
795+
DocsV3Support.LicenseRequirementChecker licenseChecker = new DocsV3Support.LicenseRequirementChecker(testClass);
796+
License.OperationMode functionLicense = licenseChecker.invoke(null);
797+
Constructor<?> ctor = constructorWithFunctionInfo(definition.clazz());
798+
if (LicenseAware.class.isAssignableFrom(definition.clazz()) == false) {
799+
// Perform simpler no-signature tests
800+
assertThat(
801+
"Function " + definition.name() + " should be licensed under " + functionLicense,
802+
functionLicense,
803+
equalTo(License.OperationMode.BASIC)
804+
);
805+
return;
806+
}
807+
// For classes with LicenseAware, we need to check that the license is correct
808+
TestCheckLicense checkLicense = new TestCheckLicense();
809+
810+
// Go through all signatures and assert that the license is as expected
811+
signatures(testClass).forEach((signature, returnType) -> {
812+
try {
813+
License.OperationMode license = licenseChecker.invoke(signature);
814+
assertNotNull("License should not be null", license);
815+
816+
// Construct an instance of the class and then call it's licenseCheck method, and compare the results
817+
Object[] args = new Object[signature.size() + 1];
818+
args[0] = Source.EMPTY;
819+
for (int i = 0; i < signature.size(); i++) {
820+
args[i + 1] = new Literal(Source.EMPTY, null, signature.get(i));
821+
}
822+
Object instance = ctor.newInstance(args);
823+
// Check that object implements the LicenseAware interface
824+
if (LicenseAware.class.isAssignableFrom(instance.getClass())) {
825+
LicenseAware licenseAware = (LicenseAware) instance;
826+
switch (license) {
827+
case BASIC -> checkLicense.assertLicenseCheck(licenseAware, signature, true, true, true);
828+
case PLATINUM -> checkLicense.assertLicenseCheck(licenseAware, signature, false, true, true);
829+
case ENTERPRISE -> checkLicense.assertLicenseCheck(licenseAware, signature, false, false, true);
830+
}
831+
} else {
832+
fail("Function " + definition.name() + " does not implement LicenseAware");
833+
}
834+
} catch (Exception e) {
835+
fail(e);
836+
}
837+
});
838+
}
839+
840+
private static class TestCheckLicense {
841+
XPackLicenseState basicLicense = makeLicenseState(License.OperationMode.BASIC);
842+
XPackLicenseState platinumLicense = makeLicenseState(License.OperationMode.PLATINUM);
843+
XPackLicenseState enterpriseLicense = makeLicenseState(License.OperationMode.ENTERPRISE);
844+
845+
private void assertLicenseCheck(
846+
LicenseAware licenseAware,
847+
List<DataType> signature,
848+
boolean allowsBasic,
849+
boolean allowsPlatinum,
850+
boolean allowsEnterprise
851+
) {
852+
boolean basic = licenseAware.licenseCheck(basicLicense);
853+
boolean platinum = licenseAware.licenseCheck(platinumLicense);
854+
boolean enterprise = licenseAware.licenseCheck(enterpriseLicense);
855+
assertThat("Basic license should be accepted for " + signature, basic, equalTo(allowsBasic));
856+
assertThat("Platinum license should be accepted for " + signature, platinum, equalTo(allowsPlatinum));
857+
assertThat("Enterprise license should be accepted for " + signature, enterprise, equalTo(allowsEnterprise));
858+
}
859+
860+
private void assertLicenseCheck(List<DataType> signature, boolean allowed, boolean expected) {
861+
assertThat("Basic license should " + (expected ? "" : "not ") + "be accepted for " + signature, allowed, equalTo(expected));
862+
}
863+
}
864+
865+
private static XPackLicenseState makeLicenseState(License.OperationMode mode) {
866+
return new XPackLicenseState(System::currentTimeMillis, new XPackLicenseStatus(mode, true, ""));
867+
}
868+
770869
/**
771870
* Asserts the result of a test case matches the expected result and warnings.
772871
* <p>
@@ -836,7 +935,7 @@ public static Map<List<DataType>, DataType> signatures(Class<?> testClass) {
836935
}
837936

838937
@AfterClass
839-
public static void renderDocs() throws IOException {
938+
public static void renderDocs() throws Exception {
840939
if (System.getProperty("generateDocs") == null) {
841940
return;
842941
}

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/DocsV3Support.java‎

Lines changed: 54 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111

1212
import org.elasticsearch.common.Strings;
1313
import org.elasticsearch.core.PathUtils;
14+
import org.elasticsearch.license.License;
1415
import org.elasticsearch.logging.LogManager;
1516
import org.elasticsearch.logging.Logger;
1617
import org.elasticsearch.xcontent.XContentBuilder;
@@ -46,6 +47,7 @@
4647
import java.io.InputStreamReader;
4748
import java.lang.annotation.Annotation;
4849
import java.lang.reflect.Constructor;
50+
import java.lang.reflect.Method;
4951
import java.nio.charset.StandardCharsets;
5052
import java.nio.file.Files;
5153
import java.nio.file.Path;
@@ -107,7 +109,7 @@ static OperatorsDocsSupport forOperators(String name, Class<?> testClass) {
107109
return new OperatorsDocsSupport(name, testClass);
108110
}
109111

110-
static void renderDocs(String name, Class<?> testClass) throws IOException {
112+
static void renderDocs(String name, Class<?> testClass) throws Exception {
111113
if (OPERATORS.containsKey(name)) {
112114
var docs = DocsV3Support.forOperators(name, testClass);
113115
docs.renderSignature();
@@ -126,7 +128,7 @@ public static void renderNegatedOperator(
126128
String name,
127129
Function<String, String> description,
128130
Class<?> testClass
129-
) throws IOException {
131+
) throws Exception {
130132
var docs = forOperators("not " + name.toLowerCase(Locale.ROOT), testClass);
131133
docs.renderDocsForNegatedOperators(ctor, description);
132134
}
@@ -272,12 +274,46 @@ public void writeToTempDir(Path dir, String extension, String str) throws IOExce
272274
}
273275
}
274276

277+
/**
278+
* This class is used to check if a license requirement method exists in the test class.
279+
* This is used to add license requirement information to the generated documentation.
280+
*/
281+
public static class LicenseRequirementChecker {
282+
private Method staticMethod;
283+
private Function<List<DataType>, License.OperationMode> fallbackLambda;
284+
285+
public LicenseRequirementChecker(Class<?> testClass) {
286+
try {
287+
staticMethod = testClass.getMethod("licenseRequirement", List.class);
288+
if (License.OperationMode.class.equals(staticMethod.getReturnType()) == false
289+
|| java.lang.reflect.Modifier.isStatic(staticMethod.getModifiers()) == false) {
290+
staticMethod = null; // Reset if the method doesn't match the signature
291+
}
292+
} catch (NoSuchMethodException e) {
293+
staticMethod = null;
294+
}
295+
296+
if (staticMethod == null) {
297+
fallbackLambda = fieldTypes -> License.OperationMode.BASIC;
298+
}
299+
}
300+
301+
public License.OperationMode invoke(List<DataType> fieldTypes) throws Exception {
302+
if (staticMethod != null) {
303+
return (License.OperationMode) staticMethod.invoke(null, fieldTypes);
304+
} else {
305+
return fallbackLambda.apply(fieldTypes);
306+
}
307+
}
308+
}
309+
275310
protected final String category;
276311
protected final String name;
277312
protected final FunctionDefinition definition;
278313
protected final Logger logger;
279314
private final Supplier<Map<List<DataType>, DataType>> signatures;
280315
private TempFileWriter tempFileWriter;
316+
private final LicenseRequirementChecker licenseChecker;
281317

282318
protected DocsV3Support(String category, String name, Class<?> testClass, Supplier<Map<List<DataType>, DataType>> signatures) {
283319
this(category, name, null, testClass, signatures);
@@ -296,6 +332,7 @@ private DocsV3Support(
296332
this.logger = LogManager.getLogger(testClass);
297333
this.signatures = signatures;
298334
this.tempFileWriter = new DocsFileWriter();
335+
this.licenseChecker = new LicenseRequirementChecker(testClass);
299336
}
300337

301338
/** Used in tests to capture output for asserting on the content */
@@ -460,7 +497,7 @@ void writeToTempKibanaDir(String subdir, String extension, String str) throws IO
460497

461498
protected abstract void renderSignature() throws IOException;
462499

463-
protected abstract void renderDocs() throws IOException;
500+
protected abstract void renderDocs() throws Exception;
464501

465502
static class FunctionDocsSupport extends DocsV3Support {
466503
private FunctionDocsSupport(String name, Class<?> testClass) {
@@ -488,7 +525,7 @@ protected void renderSignature() throws IOException {
488525
}
489526

490527
@Override
491-
protected void renderDocs() throws IOException {
528+
protected void renderDocs() throws Exception {
492529
if (definition == null) {
493530
logger.info("Skipping rendering docs because the function '{}' isn't registered", name);
494531
} else {
@@ -497,7 +534,7 @@ protected void renderDocs() throws IOException {
497534
}
498535
}
499536

500-
private void renderDocs(FunctionDefinition definition) throws IOException {
537+
private void renderDocs(FunctionDefinition definition) throws Exception {
501538
EsqlFunctionRegistry.FunctionDescription description = EsqlFunctionRegistry.description(definition);
502539
if (name.equals("case")) {
503540
/*
@@ -711,7 +748,7 @@ public void renderSignature() throws IOException {
711748
}
712749

713750
@Override
714-
public void renderDocs() throws IOException {
751+
public void renderDocs() throws Exception {
715752
Constructor<?> ctor = constructorWithFunctionInfo(op.clazz());
716753
if (ctor != null) {
717754
FunctionInfo functionInfo = ctor.getAnnotation(FunctionInfo.class);
@@ -722,7 +759,7 @@ public void renderDocs() throws IOException {
722759
}
723760
}
724761

725-
void renderDocsForNegatedOperators(Constructor<?> ctor, Function<String, String> description) throws IOException {
762+
void renderDocsForNegatedOperators(Constructor<?> ctor, Function<String, String> description) throws Exception {
726763
String baseName = name.toLowerCase(Locale.ROOT).replace("not ", "");
727764
OperatorConfig op = OPERATORS.get(baseName);
728765
assert op != null;
@@ -795,7 +832,7 @@ public Example[] examples() {
795832
}
796833

797834
void renderDocsForOperators(String name, String titleName, Constructor<?> ctor, FunctionInfo info, boolean variadic)
798-
throws IOException {
835+
throws Exception {
799836
renderKibanaInlineDocs(name, titleName, info);
800837

801838
var params = ctor.getParameters();
@@ -999,7 +1036,7 @@ void renderKibanaFunctionDefinition(
9991036
FunctionInfo info,
10001037
List<EsqlFunctionRegistry.ArgSignature> args,
10011038
boolean variadic
1002-
) throws IOException {
1039+
) throws Exception {
10031040

10041041
try (XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint().lfAtEnd().startObject()) {
10051042
builder.field(
@@ -1019,6 +1056,10 @@ void renderKibanaFunctionDefinition(
10191056
});
10201057
}
10211058
builder.field("name", name);
1059+
License.OperationMode license = licenseChecker.invoke(null);
1060+
if (license != null && license != License.OperationMode.BASIC) {
1061+
builder.field("license", license.toString());
1062+
}
10221063
if (titleName != null && titleName.equals(name) == false) {
10231064
builder.field("titleName", titleName);
10241065
}
@@ -1073,6 +1114,10 @@ void renderKibanaFunctionDefinition(
10731114
builder.endObject();
10741115
}
10751116
builder.endArray();
1117+
license = licenseChecker.invoke(sig.getKey());
1118+
if (license != null && license != License.OperationMode.BASIC) {
1119+
builder.field("license", license.toString());
1120+
}
10761121
builder.field("variadic", variadic);
10771122
builder.field("returnType", sig.getValue().esNameIfPossible());
10781123
builder.endObject();

0 commit comments

Comments
 (0)