From 9d93a1e6372c4c10410234ef1f406ba195da0631 Mon Sep 17 00:00:00 2001 From: Ignacio Vera Date: Wed, 13 Nov 2024 11:30:34 +0100 Subject: [PATCH 1/7] [ESQL] Add license check for functions --- .../core/expression/function/Function.java | 7 + .../xpack/esql/EsqlTestUtils.java | 3 +- .../xpack/esql/analysis/Verifier.java | 17 ++- .../xpack/esql/execution/PlanExecutor.java | 5 +- .../xpack/esql/plugin/EsqlPlugin.java | 9 +- .../function/CheckLicenseTests.java | 134 ++++++++++++++++++ .../LocalPhysicalPlanOptimizerTests.java | 3 +- .../esql/planner/QueryTranslatorTests.java | 3 +- .../esql/stats/PlanExecutorMetricsTests.java | 3 +- .../esql/stats/VerifierMetricsTests.java | 7 +- 10 files changed, 180 insertions(+), 11 deletions(-) create mode 100644 x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/CheckLicenseTests.java diff --git a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/function/Function.java b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/function/Function.java index cad5c631088f2..575c50a6f7231 100644 --- a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/function/Function.java +++ b/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 aggregate function can be used by a provided {@link XPackLicenseState} */ + public Predicate getLicenseChecker() { + return license -> true; + } + @Override public int hashCode() { return Objects.hash(getClass(), children()); diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java index 2913401d8aab3..d6715a932c075 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java +++ b/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() {} diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Verifier.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Verifier.java index 632f52d163349..641f4945b0002 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Verifier.java +++ b/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 failures) }); } + private void CheckLicense(LogicalPlan plan, XPackLicenseState licenseState, Set 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)) { diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/execution/PlanExecutor.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/execution/PlanExecutor.java index 816388193c5f6..c1269009c6a41 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/execution/PlanExecutor.java +++ b/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); } diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlPlugin.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlPlugin.java index e9b9f571e880e..b091ab0c1bafc 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlPlugin.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlPlugin.java @@ -38,6 +38,7 @@ 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; @@ -45,6 +46,7 @@ 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. * diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/CheckLicenseTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/CheckLicenseTests.java new file mode 100644 index 0000000000000..adf94ec859360 --- /dev/null +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/CheckLicenseTests.java @@ -0,0 +1,134 @@ +/* + * 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 dummy() | LIMIT 10"; + + public void testLicense() { + for (License.OperationMode functionLicense : License.OperationMode.values()) { + final LicensedFeature functionLicenseFeature = random().nextBoolean() + ? LicensedFeature.momentary("test", "dummy", functionLicense) + : LicensedFeature.persistent("test", "dummy", functionLicense); + final EsqlFunctionRegistry.FunctionBuilder builder = (source, expression, cfg) -> new DummyAggregateFunction( + source, + expression, + functionLicenseFeature + ); + 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 [dummy()]")); + } else { + // compliant license + assertNotNull(analyze(builder, operationMode)); + } + } + } + } + + private LogicalPlan analyze(EsqlFunctionRegistry.FunctionBuilder builder, License.OperationMode operationMode) { + final FunctionDefinition def = EsqlFunctionRegistry.def(DummyAggregateFunction.class, builder, "dummy"); + 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; + } + + private static class DummyAggregateFunction extends Function { + + private final LicensedFeature licensedFeature; + + protected DummyAggregateFunction(Source source, List children, LicensedFeature licensedFeature) { + super(source, children); + this.licensedFeature = licensedFeature; + } + + @Override + public Predicate 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 newChildren) { + throw new UnsupportedOperationException("this type of node doesn't have any children to replace"); + } + + @Override + protected NodeInfo info() { + return NodeInfo.create(this); + } + + @Override + public String getWriteableName() { + throw new UnsupportedOperationException(); + } + + @Override + public void writeTo(StreamOutput out) { + throw new UnsupportedOperationException(); + } + } + +} diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java index 073a51ee69114..2134e16b00131 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java +++ b/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)) ); } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/QueryTranslatorTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/QueryTranslatorTests.java index cf90cf96fe683..57210fda07f2b 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/QueryTranslatorTests.java +++ b/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)) ); } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/stats/PlanExecutorMetricsTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/stats/PlanExecutorMetricsTests.java index 116df21a33ac0..b323efad2b4c3 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/stats/PlanExecutorMetricsTests.java +++ b/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(); diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/stats/VerifierMetricsTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/stats/VerifierMetricsTests.java index 5e6588d2295f9..eda906b147956 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/stats/VerifierMetricsTests.java +++ b/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)); From 249d905d8ecec8f7e1a5244ea1cf9af09cdfdd53 Mon Sep 17 00:00:00 2001 From: Ignacio Vera Date: Wed, 13 Nov 2024 11:35:00 +0100 Subject: [PATCH 2/7] iter --- .../xpack/esql/core/expression/function/Function.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/function/Function.java b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/function/Function.java index 575c50a6f7231..7bc83c989c777 100644 --- a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/function/Function.java +++ b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/function/Function.java @@ -44,7 +44,7 @@ public Nullability nullable() { return Expressions.nullable(children()); } - /** Return a predicate that checks if this aggregate function can be used by a provided {@link XPackLicenseState} */ + /** Return a predicate that checks if this function can be used by a provided {@link XPackLicenseState}.*/ public Predicate getLicenseChecker() { return license -> true; } From 1c91551693604aa6ee868ccad6db29de637722a2 Mon Sep 17 00:00:00 2001 From: Ignacio Vera Date: Wed, 13 Nov 2024 12:58:47 +0100 Subject: [PATCH 3/7] iter --- .../esql/expression/function/CheckLicenseTests.java | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/CheckLicenseTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/CheckLicenseTests.java index adf94ec859360..4db60fb0faefa 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/CheckLicenseTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/CheckLicenseTests.java @@ -45,9 +45,8 @@ public void testLicense() { final LicensedFeature functionLicenseFeature = random().nextBoolean() ? LicensedFeature.momentary("test", "dummy", functionLicense) : LicensedFeature.persistent("test", "dummy", functionLicense); - final EsqlFunctionRegistry.FunctionBuilder builder = (source, expression, cfg) -> new DummyAggregateFunction( + final EsqlFunctionRegistry.FunctionBuilder builder = (source, expression, cfg) -> new DummyFunction( source, - expression, functionLicenseFeature ); for (License.OperationMode operationMode : License.OperationMode.values()) { @@ -64,7 +63,7 @@ public void testLicense() { } private LogicalPlan analyze(EsqlFunctionRegistry.FunctionBuilder builder, License.OperationMode operationMode) { - final FunctionDefinition def = EsqlFunctionRegistry.def(DummyAggregateFunction.class, builder, "dummy"); + final FunctionDefinition def = EsqlFunctionRegistry.def(DummyFunction.class, builder, "dummy"); final EsqlFunctionRegistry registry = new EsqlFunctionRegistry(def) { @Override public EsqlFunctionRegistry snapshotRegistry() { @@ -87,12 +86,12 @@ private static XPackLicenseState getLicenseState(License.OperationMode operation return licenseState; } - private static class DummyAggregateFunction extends Function { + private static class DummyFunction extends Function { private final LicensedFeature licensedFeature; - protected DummyAggregateFunction(Source source, List children, LicensedFeature licensedFeature) { - super(source, children); + private DummyFunction(Source source, LicensedFeature licensedFeature) { + super(source, List.of()); this.licensedFeature = licensedFeature; } From 6e434c56daf4357ee24a7499375d682bb9d59104 Mon Sep 17 00:00:00 2001 From: Ignacio Vera Date: Wed, 13 Nov 2024 13:15:50 +0100 Subject: [PATCH 4/7] avoid magic test failing --- .../function/CheckLicenseTests.java | 30 +++++++++++-------- 1 file changed, 18 insertions(+), 12 deletions(-) diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/CheckLicenseTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/CheckLicenseTests.java index 4db60fb0faefa..d124a1acac730 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/CheckLicenseTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/CheckLicenseTests.java @@ -38,22 +38,23 @@ public class CheckLicenseTests extends ESTestCase { private final EsqlParser parser = new EsqlParser(); - private final String esql = "from tests | eval dummy() | LIMIT 10"; + 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", "dummy", functionLicense) - : LicensedFeature.persistent("test", "dummy", functionLicense); - final EsqlFunctionRegistry.FunctionBuilder builder = (source, expression, cfg) -> new DummyFunction( - source, - functionLicenseFeature - ); + ? 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 [dummy()]")); + assertThat(ex.getMessage(), containsString("current license is non-compliant for function [license()]")); } else { // compliant license assertNotNull(analyze(builder, operationMode)); @@ -63,7 +64,7 @@ public void testLicense() { } private LogicalPlan analyze(EsqlFunctionRegistry.FunctionBuilder builder, License.OperationMode operationMode) { - final FunctionDefinition def = EsqlFunctionRegistry.def(DummyFunction.class, builder, "dummy"); + final FunctionDefinition def = EsqlFunctionRegistry.def(LicensedFunction.class, builder, "license"); final EsqlFunctionRegistry registry = new EsqlFunctionRegistry(def) { @Override public EsqlFunctionRegistry snapshotRegistry() { @@ -86,12 +87,17 @@ private static XPackLicenseState getLicenseState(License.OperationMode operation return licenseState; } - private static class DummyFunction extends Function { + // 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 final LicensedFeature licensedFeature; + private LicensedFeature licensedFeature; - private DummyFunction(Source source, LicensedFeature licensedFeature) { + public LicensedFunction(Source source) { super(source, List.of()); + } + + void setLicensedFeature(LicensedFeature licensedFeature) { this.licensedFeature = licensedFeature; } From 4852c2c10e1bfe17f78793ac7fb6537ba4ae5665 Mon Sep 17 00:00:00 2001 From: Ignacio Vera Date: Thu, 14 Nov 2024 17:14:38 +0100 Subject: [PATCH 5/7] doh! use proper case --- .../java/org/elasticsearch/xpack/esql/analysis/Verifier.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Verifier.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Verifier.java index 0bb4a9ecf9fc4..c62c7787cc959 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Verifier.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Verifier.java @@ -205,7 +205,7 @@ else if (p instanceof Lookup lookup) { checkRemoteEnrich(plan, failures); if (failures.isEmpty()) { - CheckLicense(plan, licenseState, failures); + checkLicense(plan, licenseState, failures); } // gather metrics @@ -553,7 +553,7 @@ private static void checkBinaryComparison(LogicalPlan p, Set failures) }); } - private void CheckLicense(LogicalPlan plan, XPackLicenseState licenseState, Set failures) { + private void checkLicense(LogicalPlan plan, XPackLicenseState licenseState, Set 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() + "]")); From 54ccff24a264a8389fa5f8ad21a3cd841ef038bd Mon Sep 17 00:00:00 2001 From: Ignacio Vera Date: Thu, 14 Nov 2024 17:56:21 +0100 Subject: [PATCH 6/7] Simplify method --- .../xpack/esql/core/expression/function/Function.java | 5 ++--- .../org/elasticsearch/xpack/esql/analysis/Verifier.java | 2 +- .../xpack/esql/expression/function/CheckLicenseTests.java | 7 +++---- 3 files changed, 6 insertions(+), 8 deletions(-) diff --git a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/function/Function.java b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/function/Function.java index 7bc83c989c777..e06c3e826c9f7 100644 --- a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/function/Function.java +++ b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/function/Function.java @@ -16,7 +16,6 @@ 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 @@ -45,8 +44,8 @@ public Nullability nullable() { } /** Return a predicate that checks if this function can be used by a provided {@link XPackLicenseState}.*/ - public Predicate getLicenseChecker() { - return license -> true; + public boolean checkLicense(XPackLicenseState state) { + return true; } @Override diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Verifier.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Verifier.java index c62c7787cc959..0641a03c88b69 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Verifier.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Verifier.java @@ -555,7 +555,7 @@ private static void checkBinaryComparison(LogicalPlan p, Set failures) private void checkLicense(LogicalPlan plan, XPackLicenseState licenseState, Set failures) { plan.forEachExpressionDown(Function.class, p -> { - if (p.getLicenseChecker().test(licenseState) == false) { + if (p.checkLicense(licenseState) == false) { failures.add(new Failure(p, "current license is non-compliant for function [" + p.sourceText() + "]")); } }); diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/CheckLicenseTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/CheckLicenseTests.java index d124a1acac730..98f36d339976c 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/CheckLicenseTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/CheckLicenseTests.java @@ -29,7 +29,6 @@ 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; @@ -102,11 +101,11 @@ void setLicensedFeature(LicensedFeature licensedFeature) { } @Override - public Predicate getLicenseChecker() { + public boolean checkLicense(XPackLicenseState state) { if (licensedFeature instanceof LicensedFeature.Momentary momentary) { - return momentary::check; + return momentary.check(state); } else { - return licensedFeature::checkWithoutTracking; + return licensedFeature.checkWithoutTracking(state); } } From 598c6c20f0469407081e80630fc910215b439c97 Mon Sep 17 00:00:00 2001 From: Ignacio Vera Date: Thu, 14 Nov 2024 17:58:46 +0100 Subject: [PATCH 7/7] update javadocs --- .../xpack/esql/core/expression/function/Function.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/function/Function.java b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/function/Function.java index e06c3e826c9f7..a1afcdbf1f77c 100644 --- a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/function/Function.java +++ b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/function/Function.java @@ -43,7 +43,7 @@ public Nullability nullable() { return Expressions.nullable(children()); } - /** Return a predicate that checks if this function can be used by a provided {@link XPackLicenseState}.*/ + /** Return true if this function can be executed under the provided {@link XPackLicenseState}, otherwise false.*/ public boolean checkLicense(XPackLicenseState state) { return true; }