Skip to content

Commit

Permalink
Add license check for ES|QL functions (elastic#116715) (elastic#117003)
Browse files Browse the repository at this point in the history
  • Loading branch information
iverase authored Nov 19, 2024
1 parent d8e8b6d commit 3f5251e
Show file tree
Hide file tree
Showing 10 changed files with 183 additions and 11 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -42,6 +43,11 @@ 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());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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() {}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -82,9 +83,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;
}

/**
Expand Down Expand Up @@ -201,6 +204,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);
Expand Down Expand Up @@ -546,6 +553,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.checkLicense(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)) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
);
Expand All @@ -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.
*
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,138 @@
/*
* 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 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 boolean checkLicense(XPackLicenseState state) {
if (licensedFeature instanceof LicensedFeature.Momentary momentary) {
return momentary.check(state);
} else {
return licensedFeature.checkWithoutTracking(state);
}
}

@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();
}
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
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;
Expand Down Expand Up @@ -144,7 +145,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))
);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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))
);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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));

Expand Down

0 comments on commit 3f5251e

Please sign in to comment.