Output function signature license requirements to Kibana definition files

And also test that this matches the actual licensing behaviour of the functions.

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,91 @@ 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<DataType> 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();
+        DocsV3Support.LicenseRequirementChecker licenseChecker = new DocsV3Support.LicenseRequirementChecker(testClass);
+        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");
+        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
+        XPackLicenseState basicLicense = makeLicenseState(License.OperationMode.BASIC);
+        XPackLicenseState platinumLicense = makeLicenseState(License.OperationMode.PLATINUM);
+        XPackLicenseState enterpriseLicense = makeLicenseState(License.OperationMode.ENTERPRISE);
+
+        // 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;
+                    boolean basic = licenseAware.licenseCheck(basicLicense);
+                    boolean platinum = licenseAware.licenseCheck(platinumLicense);
+                    boolean enterprise = licenseAware.licenseCheck(enterpriseLicense);
+                    if (license == License.OperationMode.BASIC) {
+                        assertTrue("Basic license should be accepted for " + signature, basic);
+                        assertTrue("Platinum license should be accepted for " + signature, platinum);
+                        assertTrue("Enterprise license should be accepted for " + signature, enterprise);
+                    }
+                    if (license == License.OperationMode.PLATINUM) {
+                        assertFalse("Basic license should not be accepted for " + signature, basic);
+                        assertTrue("Platinum license should be accepted for " + signature, platinum);
+                        assertTrue("Enterprise license should be accepted for " + signature, enterprise);
+                    }
+                    if (license == License.OperationMode.ENTERPRISE) {
+                        assertFalse("Basic license should not be accepted for " + signature, basic);
+                        assertFalse("Platinum license should not be accepted for " + signature, platinum);
+                        assertTrue("Enterprise license should be accepted for " + signature, enterprise);
+                    }
+                } else {
+                    fail("Function " + definition.name() + " does not implement LicenseAware");
+                }
+            } catch (Exception e) {
+                fail(e);
+            }
+        });
+    }
+
+    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 +954,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,49 @@ 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 -> {
+                    // Provide your default implementation here
+                    return 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 +335,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 +500,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 +528,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 +537,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 +751,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 +762,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 +835,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 +1039,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 +1059,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 +1117,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();

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

@@ -228,7 +228,7 @@ public void testRenderingExampleResultEmojis() throws IOException {
         assertThat(results, equalTo(expectedResults));
     }
 
-    public void testRenderingExampleFromClass() throws IOException {
+    public void testRenderingExampleFromClass() throws Exception {
         String expected = """
             % This is generated by ESQL's AbstractFunctionTestCase. Do no edit it. See ../README.md for how to regenerate it.
 
@@ -306,7 +306,7 @@ public void testRenderingExampleFromClass() throws IOException {
         assertThat(rendered.trim(), equalTo(expected.trim()));
     }
 
-    public void testRenderingLayoutFromClass() throws IOException {
+    public void testRenderingLayoutFromClass() throws Exception {
         String expected = """
             % This is generated by ESQL's AbstractFunctionTestCase. Do no edit it. See ../README.md for how to regenerate it.
 
@@ -353,7 +353,7 @@ public void testRenderingLayoutFromClass() throws IOException {
         assertThat(rendered.trim(), equalTo(expected.trim()));
     }
 
-    private TestDocsFileWriter renderTestClassDocs() throws IOException {
+    private TestDocsFileWriter renderTestClassDocs() throws Exception {
         FunctionInfo info = functionInfo(TestClass.class);
         assert info != null;
         FunctionDefinition definition = EsqlFunctionRegistry.def(TestClass.class, TestClass::new, "count");