Add license check for ES|QL functions

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

@@ -6,6 +6,7 @@
  */
 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;
@@ -15,6 +16,7 @@
 import java.util.Locale;
 import java.util.Objects;
 import java.util.StringJoiner;
+import java.util.function.Predicate;
 
 /**
  * Any SQL expression with parentheses, like {@code MAX()}, or {@code ABS()}. A
@@ -42,6 +44,11 @@ public Nullability nullable() {
         return Expressions.nullable(children());
     }
 
+    /** Return a predicate that checks if this function can be used by a provided {@link XPackLicenseState}.*/
+    public Predicate<XPackLicenseState> getLicenseChecker() {
+        return license -> true;
+    }
+
     @Override
     public int hashCode() {
         return Objects.hash(getClass(), children());

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java

@@ -29,6 +29,7 @@
 import org.elasticsearch.geo.GeometryTestUtils;
 import org.elasticsearch.geo.ShapeTestUtils;
 import org.elasticsearch.index.IndexMode;
+import org.elasticsearch.license.XPackLicenseState;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xcontent.json.JsonXContent;
 import org.elasticsearch.xpack.esql.action.EsqlQueryResponse;
@@ -342,7 +343,7 @@ public String toString() {
 
     public static final Configuration TEST_CFG = configuration(new QueryPragmas(Settings.EMPTY));
 
-    public static final Verifier TEST_VERIFIER = new Verifier(new Metrics(new EsqlFunctionRegistry()));
+    public static final Verifier TEST_VERIFIER = new Verifier(new Metrics(new EsqlFunctionRegistry()), new XPackLicenseState(() -> 0L));
 
     private EsqlTestUtils() {}
 

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

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.esql.analysis;
 
+import org.elasticsearch.license.XPackLicenseState;
 import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
 import org.elasticsearch.xpack.esql.common.Failure;
 import org.elasticsearch.xpack.esql.core.capabilities.Unresolvable;
@@ -78,9 +79,11 @@
 public class Verifier {
 
     private final Metrics metrics;
+    private final XPackLicenseState licenseState;
 
-    public Verifier(Metrics metrics) {
+    public Verifier(Metrics metrics, XPackLicenseState licenseState) {
         this.metrics = metrics;
+        this.licenseState = licenseState;
     }
 
     /**
@@ -197,6 +200,10 @@ else if (p instanceof Lookup lookup) {
         });
         checkRemoteEnrich(plan, failures);
 
+        if (failures.isEmpty()) {
+            CheckLicense(plan, licenseState, failures);
+        }
+
         // gather metrics
         if (failures.isEmpty()) {
             gatherMetrics(plan, partialMetrics);
@@ -473,6 +480,14 @@ 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.getLicenseChecker().test(licenseState) == false) {
+                failures.add(new Failure(p, "current license is non-compliant for function [" + p.sourceText() + "]"));
+            }
+        });
+    }
+
     private void gatherMetrics(LogicalPlan plan, BitSet b) {
         plan.forEachDown(p -> FeatureMetric.set(p, b));
         for (int i = b.nextSetBit(0); i >= 0; i = b.nextSetBit(i + 1)) {

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

@@ -9,6 +9,7 @@
 
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.indices.IndicesExpressionGrouper;
+import org.elasticsearch.license.XPackLicenseState;
 import org.elasticsearch.telemetry.metric.MeterRegistry;
 import org.elasticsearch.xpack.esql.action.EsqlExecutionInfo;
 import org.elasticsearch.xpack.esql.action.EsqlQueryRequest;
@@ -40,13 +41,13 @@ public class PlanExecutor {
     private final Verifier verifier;
     private final PlanningMetricsManager planningMetricsManager;
 
-    public PlanExecutor(IndexResolver indexResolver, MeterRegistry meterRegistry) {
+    public PlanExecutor(IndexResolver indexResolver, MeterRegistry meterRegistry, XPackLicenseState licenseState) {
         this.indexResolver = indexResolver;
         this.preAnalyzer = new PreAnalyzer();
         this.functionRegistry = new EsqlFunctionRegistry();
         this.mapper = new Mapper();
         this.metrics = new Metrics(functionRegistry);
-        this.verifier = new Verifier(metrics);
+        this.verifier = new Verifier(metrics, licenseState);
         this.planningMetricsManager = new PlanningMetricsManager(meterRegistry);
     }
 

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

@@ -38,13 +38,15 @@
 import org.elasticsearch.compute.operator.exchange.ExchangeSourceOperator;
 import org.elasticsearch.compute.operator.topn.TopNOperatorStatus;
 import org.elasticsearch.features.NodeFeature;
+import org.elasticsearch.license.XPackLicenseState;
 import org.elasticsearch.plugins.ActionPlugin;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.rest.RestController;
 import org.elasticsearch.rest.RestHandler;
 import org.elasticsearch.threadpool.ExecutorBuilder;
 import org.elasticsearch.threadpool.FixedExecutorBuilder;
 import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.xpack.core.XPackPlugin;
 import org.elasticsearch.xpack.core.action.XPackInfoFeatureAction;
 import org.elasticsearch.xpack.core.action.XPackUsageFeatureAction;
 import org.elasticsearch.xpack.esql.EsqlInfoTransportAction;
@@ -116,7 +118,7 @@ public Collection<?> createComponents(PluginServices services) {
         BlockFactory blockFactory = new BlockFactory(circuitBreaker, bigArrays, maxPrimitiveArrayBlockSize);
         setupSharedSecrets();
         return List.of(
-            new PlanExecutor(new IndexResolver(services.client()), services.telemetryProvider().getMeterRegistry()),
+            new PlanExecutor(new IndexResolver(services.client()), services.telemetryProvider().getMeterRegistry(), getLicenseState()),
             new ExchangeService(services.clusterService().getSettings(), services.threadPool(), ThreadPool.Names.SEARCH, blockFactory),
             blockFactory
         );
@@ -131,6 +133,11 @@ private void setupSharedSecrets() {
         }
     }
 
+    // to be overriden by tests
+    protected XPackLicenseState getLicenseState() {
+        return XPackPlugin.getSharedLicenseState();
+    }
+
     /**
      * The settings defined by the ESQL plugin.
      *

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

@@ -0,0 +1,139 @@
+/*
+ * 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.expression.function;
+
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.license.License;
+import org.elasticsearch.license.LicensedFeature;
+import org.elasticsearch.license.TestUtils;
+import org.elasticsearch.license.XPackLicenseState;
+import org.elasticsearch.license.internal.XPackLicenseStatus;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.esql.EsqlTestUtils;
+import org.elasticsearch.xpack.esql.VerificationException;
+import org.elasticsearch.xpack.esql.analysis.Analyzer;
+import org.elasticsearch.xpack.esql.analysis.AnalyzerContext;
+import org.elasticsearch.xpack.esql.analysis.Verifier;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.function.Function;
+import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
+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.LogicalPlan;
+import org.elasticsearch.xpack.esql.stats.Metrics;
+
+import java.util.List;
+import java.util.function.Predicate;
+
+import static org.elasticsearch.xpack.esql.analysis.AnalyzerTestUtils.analyzerDefaultMapping;
+import static org.elasticsearch.xpack.esql.analysis.AnalyzerTestUtils.defaultEnrichResolution;
+import static org.hamcrest.Matchers.containsString;
+
+public class CheckLicenseTests extends ESTestCase {
+
+    private final EsqlParser parser = new EsqlParser();
+    private final String esql = "from tests | eval license() | LIMIT 10";
+
+    public void testLicense() {
+        for (License.OperationMode functionLicense : License.OperationMode.values()) {
+            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()]"));
+                } else {
+                    // compliant license
+                    assertNotNull(analyze(builder, operationMode));
+                }
+            }
+        }
+    }
+
+    private LogicalPlan analyze(EsqlFunctionRegistry.FunctionBuilder builder, License.OperationMode operationMode) {
+        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));
+    }
+
+    private static Analyzer analyzer(EsqlFunctionRegistry registry, License.OperationMode operationMode) {
+        return new Analyzer(
+            new AnalyzerContext(EsqlTestUtils.TEST_CFG, registry, analyzerDefaultMapping(), defaultEnrichResolution()),
+            new Verifier(new Metrics(new EsqlFunctionRegistry()), getLicenseState(operationMode))
+        );
+    }
+
+    private static XPackLicenseState getLicenseState(License.OperationMode operationMode) {
+        final TestUtils.UpdatableLicenseState licenseState = new TestUtils.UpdatableLicenseState();
+        licenseState.update(new XPackLicenseStatus(operationMode, true, null));
+        return licenseState;
+    }
+
+    // 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 {
+
+        private LicensedFeature licensedFeature;
+
+        public LicensedFunction(Source source) {
+            super(source, List.of());
+        }
+
+        void setLicensedFeature(LicensedFeature licensedFeature) {
+            this.licensedFeature = licensedFeature;
+        }
+
+        @Override
+        public Predicate<XPackLicenseState> getLicenseChecker() {
+            if (licensedFeature instanceof LicensedFeature.Momentary momentary) {
+                return momentary::check;
+            } else {
+                return licensedFeature::checkWithoutTracking;
+            }
+        }
+
+        @Override
+        public DataType dataType() {
+            return DataType.KEYWORD;
+        }
+
+        @Override
+        public Expression replaceChildren(List<Expression> newChildren) {
+            throw new UnsupportedOperationException("this type of node doesn't have any children to replace");
+        }
+
+        @Override
+        protected NodeInfo<? extends Expression> info() {
+            return NodeInfo.create(this);
+        }
+
+        @Override
+        public String getWriteableName() {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) {
+            throw new UnsupportedOperationException();
+        }
+    }
+
+}

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java

@@ -22,6 +22,7 @@
 import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.index.query.RangeQueryBuilder;
 import org.elasticsearch.index.query.SearchExecutionContext;
+import org.elasticsearch.license.XPackLicenseState;
 import org.elasticsearch.xpack.core.enrich.EnrichPolicy;
 import org.elasticsearch.xpack.esql.EsqlTestUtils;
 import org.elasticsearch.xpack.esql.EsqlTestUtils.TestSearchStats;
@@ -145,7 +146,7 @@ private Analyzer makeAnalyzer(String mappingFileName, EnrichResolution enrichRes
 
         return new Analyzer(
             new AnalyzerContext(config, new EsqlFunctionRegistry(), getIndexResult, enrichResolution),
-            new Verifier(new Metrics(new EsqlFunctionRegistry()))
+            new Verifier(new Metrics(new EsqlFunctionRegistry()), new XPackLicenseState(() -> 0L))
         );
     }
 

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/QueryTranslatorTests.java

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.esql.planner;
 
 import org.elasticsearch.index.IndexMode;
+import org.elasticsearch.license.XPackLicenseState;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xpack.esql.EsqlTestUtils;
 import org.elasticsearch.xpack.esql.analysis.Analyzer;
@@ -46,7 +47,7 @@ private static Analyzer makeAnalyzer(String mappingFileName) {
 
         return new Analyzer(
             new AnalyzerContext(EsqlTestUtils.TEST_CFG, new EsqlFunctionRegistry(), getIndexResult, new EnrichResolution()),
-            new Verifier(new Metrics(new EsqlFunctionRegistry()))
+            new Verifier(new Metrics(new EsqlFunctionRegistry()), new XPackLicenseState(() -> 0L))
         );
     }
 

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/stats/PlanExecutorMetricsTests.java

@@ -17,6 +17,7 @@
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.index.IndexMode;
 import org.elasticsearch.indices.IndicesExpressionGrouper;
+import org.elasticsearch.license.XPackLicenseState;
 import org.elasticsearch.telemetry.metric.MeterRegistry;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.threadpool.TestThreadPool;
@@ -102,7 +103,7 @@ public void testFailedMetric() {
             return null;
         }).when(esqlClient).execute(eq(EsqlResolveFieldsAction.TYPE), any(), any());
 
-        var planExecutor = new PlanExecutor(indexResolver, MeterRegistry.NOOP);
+        var planExecutor = new PlanExecutor(indexResolver, MeterRegistry.NOOP, new XPackLicenseState(() -> 0L));
         var enrichResolver = mockEnrichResolver();
 
         var request = new EsqlQueryRequest();

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/stats/VerifierMetricsTests.java

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.esql.stats;
 
+import org.elasticsearch.license.XPackLicenseState;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xpack.core.watcher.common.stats.Counters;
 import org.elasticsearch.xpack.esql.analysis.Verifier;
@@ -205,7 +206,7 @@ public void testTwoWhereQuery() {
 
     public void testTwoQueriesExecuted() {
         Metrics metrics = new Metrics(new EsqlFunctionRegistry());
-        Verifier verifier = new Verifier(metrics);
+        Verifier verifier = new Verifier(metrics, new XPackLicenseState(() -> 0L));
         esqlWithVerifier("""
                from employees
                | where languages > 2
@@ -252,7 +253,7 @@ public void testTwoQueriesExecuted() {
 
     public void testMultipleFunctions() {
         Metrics metrics = new Metrics(new EsqlFunctionRegistry());
-        Verifier verifier = new Verifier(metrics);
+        Verifier verifier = new Verifier(metrics, new XPackLicenseState(() -> 0L));
         esqlWithVerifier("""
                from employees
                | where languages > 2
@@ -526,7 +527,7 @@ private Counters esql(String esql, Verifier v) {
         Metrics metrics = null;
         if (v == null) {
             metrics = new Metrics(new EsqlFunctionRegistry());
-            verifier = new Verifier(metrics);
+            verifier = new Verifier(metrics, new XPackLicenseState(() -> 0L));
         }
         analyzer(verifier).analyze(parser.createStatement(esql));