ESQL: Add a LicenseAware interface for licensed Nodes (#118931)

This adds a new interface that elements that require a proper license state can implement to enforce the license requirement. This can be now applied to any node or node property.

The check still happens in the Verifier, since the plan needs to be analysed first and the check still only happens if no other verification faults exist already.

Fixes #117405

Author: bpintea

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);
+    }
 }