ESQL: Add a LicenseAware interface for licensed Nodes

docs/changelog/118931.yaml

@@ -0,0 +1,6 @@
+pr: 118931
+summary: Add a `LicenseAware` interface for licensed Nodes
+area: ES|QL
+type: enhancement
+issues:
+ - 117405

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/function/Function.java

@@ -6,7 +6,6 @@
  */
 package org.elasticsearch.xpack.esql.core.expression.function;
 
-import org.elasticsearch.license.XPackLicenseState;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
 import org.elasticsearch.xpack.esql.core.expression.Nullability;
@@ -43,11 +42,6 @@ public Nullability nullable() {
         return Expressions.nullable(children());
     }
 
-    /** Return true if this function can be executed under the provided {@link XPackLicenseState}, otherwise false.*/
-    public boolean checkLicense(XPackLicenseState state) {
-        return true;
-    }
-
     @Override
     public int hashCode() {
         return Objects.hash(getClass(), children());

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/LicenseAware.java

@@ -0,0 +1,15 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql;
+
+import org.elasticsearch.license.XPackLicenseState;
+
+public interface LicenseAware {
+    /** Return true if the implementer can be executed under the provided {@link XPackLicenseState}, otherwise false.*/
+    boolean licenseCheck(XPackLicenseState state);
+}

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Verifier.java

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.esql.analysis;
 
 import org.elasticsearch.license.XPackLicenseState;
+import org.elasticsearch.xpack.esql.LicenseAware;
 import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
 import org.elasticsearch.xpack.esql.common.Failure;
 import org.elasticsearch.xpack.esql.core.capabilities.Unresolvable;
@@ -26,6 +27,7 @@
 import org.elasticsearch.xpack.esql.core.expression.predicate.logical.Not;
 import org.elasticsearch.xpack.esql.core.expression.predicate.logical.Or;
 import org.elasticsearch.xpack.esql.core.expression.predicate.operator.comparison.BinaryComparison;
+import org.elasticsearch.xpack.esql.core.tree.Node;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.expression.function.UnsupportedAttribute;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.AggregateFunction;
@@ -209,7 +211,7 @@ else if (p instanceof Lookup lookup) {
         checkRemoteEnrich(plan, failures);
 
         if (failures.isEmpty()) {
-            checkLicense(plan, licenseState, failures);
+            licenseCheck(plan, failures);
         }
 
         // gather metrics
@@ -587,11 +589,15 @@ private static void checkBinaryComparison(LogicalPlan p, Set<Failure> failures)
         });
     }
 
-    private void checkLicense(LogicalPlan plan, XPackLicenseState licenseState, Set<Failure> failures) {
-        plan.forEachExpressionDown(Function.class, p -> {
-            if (p.checkLicense(licenseState) == false) {
-                failures.add(new Failure(p, "current license is non-compliant for function [" + p.sourceText() + "]"));
+    private void licenseCheck(LogicalPlan plan, Set<Failure> failures) {
+        Consumer<Node<?>> licenseCheck = n -> {
+            if (n instanceof LicenseAware la && la.licenseCheck(licenseState) == false) {
+                failures.add(fail(n, "current license is non-compliant for [{}]", n.sourceText()));
             }
+        };
+        plan.forEachDown(p -> {
+            licenseCheck.accept(p);
+            p.forEachExpression(Expression.class, licenseCheck);
         });
     }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialAggregateFunction.java

@@ -11,6 +11,7 @@
 import org.elasticsearch.index.mapper.MappedFieldType.FieldExtractPreference;
 import org.elasticsearch.license.License;
 import org.elasticsearch.license.XPackLicenseState;
+import org.elasticsearch.xpack.esql.LicenseAware;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 
@@ -24,7 +25,7 @@
  * The AggregateMapper class will generate multiple aggregation functions for each combination, allowing the planner to
  * select the best one.
  */
-public abstract class SpatialAggregateFunction extends AggregateFunction {
+public abstract class SpatialAggregateFunction extends AggregateFunction implements LicenseAware {
     protected final FieldExtractPreference fieldExtractPreference;
 
     protected SpatialAggregateFunction(Source source, Expression field, Expression filter, FieldExtractPreference fieldExtractPreference) {
@@ -41,7 +42,7 @@ protected SpatialAggregateFunction(StreamInput in, FieldExtractPreference fieldE
     public abstract SpatialAggregateFunction withDocValues();
 
     @Override
-    public boolean checkLicense(XPackLicenseState state) {
+    public boolean licenseCheck(XPackLicenseState state) {
         return switch (field().dataType()) {
             case GEO_SHAPE, CARTESIAN_SHAPE -> state.isAllowedByLicense(License.OperationMode.PLATINUM);
             default -> true;

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

@@ -15,6 +15,7 @@
 import org.elasticsearch.license.internal.XPackLicenseStatus;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xpack.esql.EsqlTestUtils;
+import org.elasticsearch.xpack.esql.LicenseAware;
 import org.elasticsearch.xpack.esql.VerificationException;
 import org.elasticsearch.xpack.esql.analysis.Analyzer;
 import org.elasticsearch.xpack.esql.analysis.AnalyzerContext;
@@ -25,10 +26,12 @@
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.parser.EsqlParser;
+import org.elasticsearch.xpack.esql.plan.logical.Limit;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.stats.Metrics;
 
 import java.util.List;
+import java.util.Objects;
 
 import static org.elasticsearch.xpack.esql.analysis.AnalyzerTestUtils.analyzerDefaultMapping;
 import static org.elasticsearch.xpack.esql.analysis.AnalyzerTestUtils.defaultEnrichResolution;
@@ -44,33 +47,42 @@ public void testLicense() {
             final LicensedFeature functionLicenseFeature = random().nextBoolean()
                 ? LicensedFeature.momentary("test", "license", functionLicense)
                 : LicensedFeature.persistent("test", "license", functionLicense);
-            final EsqlFunctionRegistry.FunctionBuilder builder = (source, expression, cfg) -> {
-                final LicensedFunction licensedFunction = new LicensedFunction(source);
-                licensedFunction.setLicensedFeature(functionLicenseFeature);
-                return licensedFunction;
-            };
             for (License.OperationMode operationMode : License.OperationMode.values()) {
                 if (License.OperationMode.TRIAL != operationMode && License.OperationMode.compare(operationMode, functionLicense) < 0) {
                     // non-compliant license
-                    final VerificationException ex = expectThrows(VerificationException.class, () -> analyze(builder, operationMode));
-                    assertThat(ex.getMessage(), containsString("current license is non-compliant for function [license()]"));
+                    final VerificationException ex = expectThrows(
+                        VerificationException.class,
+                        () -> analyze(operationMode, functionLicenseFeature)
+                    );
+                    assertThat(ex.getMessage(), containsString("current license is non-compliant for [license()]"));
+                    assertThat(ex.getMessage(), containsString("current license is non-compliant for [LicensedLimit]"));
                 } else {
                     // compliant license
-                    assertNotNull(analyze(builder, operationMode));
+                    assertNotNull(analyze(operationMode, functionLicenseFeature));
                 }
             }
         }
     }
 
-    private LogicalPlan analyze(EsqlFunctionRegistry.FunctionBuilder builder, License.OperationMode operationMode) {
+    private LogicalPlan analyze(License.OperationMode operationMode, LicensedFeature functionLicenseFeature) {
+        final EsqlFunctionRegistry.FunctionBuilder builder = (source, expression, cfg) -> new LicensedFunction(
+            source,
+            functionLicenseFeature
+        );
         final FunctionDefinition def = EsqlFunctionRegistry.def(LicensedFunction.class, builder, "license");
         final EsqlFunctionRegistry registry = new EsqlFunctionRegistry(def) {
             @Override
             public EsqlFunctionRegistry snapshotRegistry() {
                 return this;
             }
         };
-        return analyzer(registry, operationMode).analyze(parser.createStatement(esql));
+
+        var plan = parser.createStatement(esql);
+        plan = plan.transformDown(
+            Limit.class,
+            l -> Objects.equals(l.limit().fold(), 10) ? new LicensedLimit(l.source(), l.limit(), l.child(), functionLicenseFeature) : l
+        );
+        return analyzer(registry, operationMode).analyze(plan);
     }
 
     private static Analyzer analyzer(EsqlFunctionRegistry registry, License.OperationMode operationMode) {
@@ -88,25 +100,18 @@ private static XPackLicenseState getLicenseState(License.OperationMode operation
 
     // It needs to be public because we run validation on it via reflection in org.elasticsearch.xpack.esql.tree.EsqlNodeSubclassTests.
     // This test prevents to add the license as constructor parameter too.
-    public static class LicensedFunction extends Function {
+    public static class LicensedFunction extends Function implements LicenseAware {
 
-        private LicensedFeature licensedFeature;
+        private final LicensedFeature licensedFeature;
 
-        public LicensedFunction(Source source) {
+        public LicensedFunction(Source source, LicensedFeature licensedFeature) {
             super(source, List.of());
-        }
-
-        void setLicensedFeature(LicensedFeature licensedFeature) {
             this.licensedFeature = licensedFeature;
         }
 
         @Override
-        public boolean checkLicense(XPackLicenseState state) {
-            if (licensedFeature instanceof LicensedFeature.Momentary momentary) {
-                return momentary.check(state);
-            } else {
-                return licensedFeature.checkWithoutTracking(state);
-            }
+        public boolean licenseCheck(XPackLicenseState state) {
+            return checkLicense(state, licensedFeature);
         }
 
         @Override
@@ -121,7 +126,7 @@ public Expression replaceChildren(List<Expression> newChildren) {
 
         @Override
         protected NodeInfo<? extends Expression> info() {
-            return NodeInfo.create(this);
+            return NodeInfo.create(this, LicensedFunction::new, licensedFeature);
         }
 
         @Override
@@ -135,4 +140,39 @@ public void writeTo(StreamOutput out) {
         }
     }
 
+    public static class LicensedLimit extends Limit implements LicenseAware {
+
+        private final LicensedFeature licensedFeature;
+
+        public LicensedLimit(Source source, Expression limit, LogicalPlan child, LicensedFeature licensedFeature) {
+            super(source, limit, child);
+            this.licensedFeature = licensedFeature;
+        }
+
+        @Override
+        public boolean licenseCheck(XPackLicenseState state) {
+            return checkLicense(state, licensedFeature);
+        }
+
+        @Override
+        public Limit replaceChild(LogicalPlan newChild) {
+            return new LicensedLimit(source(), limit(), newChild, licensedFeature);
+        }
+
+        @Override
+        protected NodeInfo<Limit> info() {
+            return NodeInfo.create(this, LicensedLimit::new, limit(), child(), licensedFeature);
+        }
+
+        @Override
+        public String sourceText() {
+            return "LicensedLimit";
+        }
+    }
+
+    private static boolean checkLicense(XPackLicenseState state, LicensedFeature licensedFeature) {
+        return licensedFeature instanceof LicensedFeature.Momentary momentary
+            ? momentary.check(state)
+            : licensedFeature.checkWithoutTracking(state);
+    }
 }