Output function signature license requirements to Kibana definitions (elastic#127717)

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.

Author: craigtaverner

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

@@ -560,6 +560,10 @@ public static boolean isSpatialPoint(DataType t) {
         return t == GEO_POINT || t == CARTESIAN_POINT;
     }
 
+    public static boolean isSpatialShape(DataType t) {
+        return t == GEO_SHAPE || t == CARTESIAN_SHAPE;
+    }
+
     public static boolean isSpatialGeo(DataType t) {
         return t == GEO_POINT || t == GEO_SHAPE;
     }

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

@@ -26,9 +26,13 @@
 import org.elasticsearch.compute.test.BlockTestUtils;
 import org.elasticsearch.compute.test.TestBlockFactory;
 import org.elasticsearch.indices.CrankyCircuitBreakerService;
+import org.elasticsearch.license.License;
+import org.elasticsearch.license.XPackLicenseState;
+import org.elasticsearch.license.internal.XPackLicenseStatus;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
 import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.esql.LicenseAware;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
@@ -54,7 +58,6 @@
 import org.junit.After;
 import org.junit.AfterClass;
 
-import java.io.IOException;
 import java.lang.reflect.Constructor;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
@@ -730,7 +733,8 @@ public void testSerializationOfSimple() {
      */
     @AfterClass
     public static void testFunctionInfo() {
-        Logger log = LogManager.getLogger(getTestClass());
+        Class<?> testClass = getTestClass();
+        Logger log = LogManager.getLogger(testClass);
         FunctionDefinition definition = definition(functionName());
         if (definition == null) {
             log.info("Skipping function info checks because the function isn't registered");
@@ -753,7 +757,7 @@ public static void testFunctionInfo() {
         for (int i = 0; i < args.size(); i++) {
             typesFromSignature.add(new HashSet<>());
         }
-        for (Map.Entry<List<DataType>, DataType> entry : signatures(getTestClass()).entrySet()) {
+        for (Map.Entry<List<DataType>, DataType> entry : signatures(testClass).entrySet()) {
             List<DataType> types = entry.getKey();
             for (int i = 0; i < args.size() && i < types.size(); i++) {
                 typesFromSignature.get(i).add(types.get(i).esNameIfPossible());
@@ -796,6 +800,101 @@ public static void testFunctionInfo() {
         assertEquals(returnFromSignature, returnTypes);
     }
 
+    /**
+     * This test is meant to validate that the license checks documented match those enforced.
+     * The expectations are set in the test class using a method with this signature:
+     * <code>
+     *     public static License.OperationMode licenseRequirement(List&lt;DataType&gt; fieldTypes);
+     * </code>
+     * License enforcement in the function class is achieved using the interface <code>LicenseAware</code>.
+     * This test will make sure the two are in agreement, and does not require that the function class actually
+     * report its license level. If we add license checks to any function, but fail to also add the expected
+     * license level to the test class, this test will fail.
+     */
+    @AfterClass
+    public static void testFunctionLicenseChecks() throws Exception {
+        Class<?> testClass = getTestClass();
+        Logger log = LogManager.getLogger(testClass);
+        FunctionDefinition definition = definition(functionName());
+        if (definition == null) {
+            log.info("Skipping function info checks because the function isn't registered");
+            return;
+        }
+        log.info("Running function license checks");
+        DocsV3Support.LicenseRequirementChecker licenseChecker = new DocsV3Support.LicenseRequirementChecker(testClass);
+        License.OperationMode functionLicense = licenseChecker.invoke(null);
+        Constructor<?> ctor = constructorWithFunctionInfo(definition.clazz());
+        if (LicenseAware.class.isAssignableFrom(definition.clazz()) == false) {
+            // Perform simpler no-signature tests
+            assertThat(
+                "Function " + definition.name() + " should be licensed under " + functionLicense,
+                functionLicense,
+                equalTo(License.OperationMode.BASIC)
+            );
+            return;
+        }
+        // For classes with LicenseAware, we need to check that the license is correct
+        TestCheckLicense checkLicense = new TestCheckLicense();
+
+        // Go through all signatures and assert that the license is as expected
+        signatures(testClass).forEach((signature, returnType) -> {
+            try {
+                License.OperationMode license = licenseChecker.invoke(signature);
+                assertNotNull("License should not be null", license);
+
+                // Construct an instance of the class and then call it's licenseCheck method, and compare the results
+                Object[] args = new Object[signature.size() + 1];
+                args[0] = Source.EMPTY;
+                for (int i = 0; i < signature.size(); i++) {
+                    args[i + 1] = new Literal(Source.EMPTY, null, signature.get(i));
+                }
+                Object instance = ctor.newInstance(args);
+                // Check that object implements the LicenseAware interface
+                if (LicenseAware.class.isAssignableFrom(instance.getClass())) {
+                    LicenseAware licenseAware = (LicenseAware) instance;
+                    switch (license) {
+                        case BASIC -> checkLicense.assertLicenseCheck(licenseAware, signature, true, true, true);
+                        case PLATINUM -> checkLicense.assertLicenseCheck(licenseAware, signature, false, true, true);
+                        case ENTERPRISE -> checkLicense.assertLicenseCheck(licenseAware, signature, false, false, true);
+                    }
+                } else {
+                    fail("Function " + definition.name() + " does not implement LicenseAware");
+                }
+            } catch (Exception e) {
+                fail(e);
+            }
+        });
+    }
+
+    private static class TestCheckLicense {
+        XPackLicenseState basicLicense = makeLicenseState(License.OperationMode.BASIC);
+        XPackLicenseState platinumLicense = makeLicenseState(License.OperationMode.PLATINUM);
+        XPackLicenseState enterpriseLicense = makeLicenseState(License.OperationMode.ENTERPRISE);
+
+        private void assertLicenseCheck(
+            LicenseAware licenseAware,
+            List<DataType> signature,
+            boolean allowsBasic,
+            boolean allowsPlatinum,
+            boolean allowsEnterprise
+        ) {
+            boolean basic = licenseAware.licenseCheck(basicLicense);
+            boolean platinum = licenseAware.licenseCheck(platinumLicense);
+            boolean enterprise = licenseAware.licenseCheck(enterpriseLicense);
+            assertThat("Basic license should be accepted for " + signature, basic, equalTo(allowsBasic));
+            assertThat("Platinum license should be accepted for " + signature, platinum, equalTo(allowsPlatinum));
+            assertThat("Enterprise license should be accepted for " + signature, enterprise, equalTo(allowsEnterprise));
+        }
+
+        private void assertLicenseCheck(List<DataType> signature, boolean allowed, boolean expected) {
+            assertThat("Basic license should " + (expected ? "" : "not ") + "be accepted for " + signature, allowed, equalTo(expected));
+        }
+    }
+
+    private static XPackLicenseState makeLicenseState(License.OperationMode mode) {
+        return new XPackLicenseState(System::currentTimeMillis, new XPackLicenseStatus(mode, true, ""));
+    }
+
     /**
      * Asserts the result of a test case matches the expected result and warnings.
      * <p>
@@ -865,7 +964,7 @@ public static Map<List<DataType>, DataType> signatures(Class<?> testClass) {
     }
 
     @AfterClass
-    public static void renderDocs() throws IOException {
+    public static void renderDocs() throws Exception {
         if (System.getProperty("generateDocs") == null) {
             return;
         }

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

@@ -11,6 +11,7 @@
 
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.core.PathUtils;
+import org.elasticsearch.license.License;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
 import org.elasticsearch.xcontent.XContentBuilder;
@@ -46,6 +47,7 @@
 import java.io.InputStreamReader;
 import java.lang.annotation.Annotation;
 import java.lang.reflect.Constructor;
+import java.lang.reflect.Method;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -107,7 +109,7 @@ static OperatorsDocsSupport forOperators(String name, Class<?> testClass) {
         return new OperatorsDocsSupport(name, testClass);
     }
 
-    static void renderDocs(String name, Class<?> testClass) throws IOException {
+    static void renderDocs(String name, Class<?> testClass) throws Exception {
         if (OPERATORS.containsKey(name)) {
             var docs = DocsV3Support.forOperators(name, testClass);
             docs.renderSignature();
@@ -126,7 +128,7 @@ public static void renderNegatedOperator(
         String name,
         Function<String, String> description,
         Class<?> testClass
-    ) throws IOException {
+    ) throws Exception {
         var docs = forOperators("not " + name.toLowerCase(Locale.ROOT), testClass);
         docs.renderDocsForNegatedOperators(ctor, description);
     }
@@ -272,12 +274,46 @@ public void writeToTempDir(Path dir, String extension, String str) throws IOExce
         }
     }
 
+    /**
+     * This class is used to check if a license requirement method exists in the test class.
+     * This is used to add license requirement information to the generated documentation.
+     */
+    public static class LicenseRequirementChecker {
+        private Method staticMethod;
+        private Function<List<DataType>, License.OperationMode> fallbackLambda;
+
+        public LicenseRequirementChecker(Class<?> testClass) {
+            try {
+                staticMethod = testClass.getMethod("licenseRequirement", List.class);
+                if (License.OperationMode.class.equals(staticMethod.getReturnType()) == false
+                    || java.lang.reflect.Modifier.isStatic(staticMethod.getModifiers()) == false) {
+                    staticMethod = null; // Reset if the method doesn't match the signature
+                }
+            } catch (NoSuchMethodException e) {
+                staticMethod = null;
+            }
+
+            if (staticMethod == null) {
+                fallbackLambda = fieldTypes -> License.OperationMode.BASIC;
+            }
+        }
+
+        public License.OperationMode invoke(List<DataType> fieldTypes) throws Exception {
+            if (staticMethod != null) {
+                return (License.OperationMode) staticMethod.invoke(null, fieldTypes);
+            } else {
+                return fallbackLambda.apply(fieldTypes);
+            }
+        }
+    }
+
     protected final String category;
     protected final String name;
     protected final FunctionDefinition definition;
     protected final Logger logger;
     private final Supplier<Map<List<DataType>, DataType>> signatures;
     private TempFileWriter tempFileWriter;
+    private final LicenseRequirementChecker licenseChecker;
 
     protected DocsV3Support(String category, String name, Class<?> testClass, Supplier<Map<List<DataType>, DataType>> signatures) {
         this(category, name, null, testClass, signatures);
@@ -296,6 +332,7 @@ private DocsV3Support(
         this.logger = LogManager.getLogger(testClass);
         this.signatures = signatures;
         this.tempFileWriter = new DocsFileWriter();
+        this.licenseChecker = new LicenseRequirementChecker(testClass);
     }
 
     /** 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
 
     protected abstract void renderSignature() throws IOException;
 
-    protected abstract void renderDocs() throws IOException;
+    protected abstract void renderDocs() throws Exception;
 
     static class FunctionDocsSupport extends DocsV3Support {
         private FunctionDocsSupport(String name, Class<?> testClass) {
@@ -488,7 +525,7 @@ protected void renderSignature() throws IOException {
         }
 
         @Override
-        protected void renderDocs() throws IOException {
+        protected void renderDocs() throws Exception {
             if (definition == null) {
                 logger.info("Skipping rendering docs because the function '{}' isn't registered", name);
             } else {
@@ -497,7 +534,7 @@ protected void renderDocs() throws IOException {
             }
         }
 
-        private void renderDocs(FunctionDefinition definition) throws IOException {
+        private void renderDocs(FunctionDefinition definition) throws Exception {
             EsqlFunctionRegistry.FunctionDescription description = EsqlFunctionRegistry.description(definition);
             if (name.equals("case")) {
                 /*
@@ -711,7 +748,7 @@ public void renderSignature() throws IOException {
         }
 
         @Override
-        public void renderDocs() throws IOException {
+        public void renderDocs() throws Exception {
             Constructor<?> ctor = constructorWithFunctionInfo(op.clazz());
             if (ctor != null) {
                 FunctionInfo functionInfo = ctor.getAnnotation(FunctionInfo.class);
@@ -722,7 +759,7 @@ public void renderDocs() throws IOException {
             }
         }
 
-        void renderDocsForNegatedOperators(Constructor<?> ctor, Function<String, String> description) throws IOException {
+        void renderDocsForNegatedOperators(Constructor<?> ctor, Function<String, String> description) throws Exception {
             String baseName = name.toLowerCase(Locale.ROOT).replace("not ", "");
             OperatorConfig op = OPERATORS.get(baseName);
             assert op != null;
@@ -795,7 +832,7 @@ public Example[] examples() {
         }
 
         void renderDocsForOperators(String name, String titleName, Constructor<?> ctor, FunctionInfo info, boolean variadic)
-            throws IOException {
+            throws Exception {
             renderKibanaInlineDocs(name, titleName, info);
 
             var params = ctor.getParameters();
@@ -999,7 +1036,7 @@ void renderKibanaFunctionDefinition(
         FunctionInfo info,
         List<EsqlFunctionRegistry.ArgSignature> args,
         boolean variadic
-    ) throws IOException {
+    ) throws Exception {
 
         try (XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint().lfAtEnd().startObject()) {
             builder.field(
@@ -1019,6 +1056,10 @@ void renderKibanaFunctionDefinition(
                 });
             }
             builder.field("name", name);
+            License.OperationMode license = licenseChecker.invoke(null);
+            if (license != null && license != License.OperationMode.BASIC) {
+                builder.field("license", license.toString());
+            }
             if (titleName != null && titleName.equals(name) == false) {
                 builder.field("titleName", titleName);
             }
@@ -1073,6 +1114,10 @@ void renderKibanaFunctionDefinition(
                         builder.endObject();
                     }
                     builder.endArray();
+                    license = licenseChecker.invoke(sig.getKey());
+                    if (license != null && license != License.OperationMode.BASIC) {
+                        builder.field("license", license.toString());
+                    }
                     builder.field("variadic", variadic);
                     builder.field("returnType", sig.getValue().esNameIfPossible());
                     builder.endObject();