Skip to content

Suppport per-project behavior in ESQL extra verifiers #131884

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -87,10 +87,6 @@ protected Settings nodeSettings(int nodeOrdinal, Settings otherSettings) {
}

public static class EsqlTestPluginWithMockBlockFactory extends EsqlPlugin {
public EsqlTestPluginWithMockBlockFactory(Settings settings) {
super(settings);
}

@Override
protected BlockFactoryProvider blockFactoryProvider(
CircuitBreaker breaker,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@

package org.elasticsearch.xpack.esql.action;

import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.license.License;
import org.elasticsearch.license.XPackLicenseState;
import org.elasticsearch.license.internal.XPackLicenseStatus;
Expand All @@ -20,10 +19,6 @@
* that require an Enteprise (or Trial) license.
*/
public class EsqlPluginWithEnterpriseOrTrialLicense extends EsqlPlugin {
public EsqlPluginWithEnterpriseOrTrialLicense(Settings settings) {
super(settings);
}

protected XPackLicenseState getLicenseState() {
License.OperationMode operationMode = randomFrom(License.OperationMode.ENTERPRISE, License.OperationMode.TRIAL);
return new XPackLicenseState(() -> System.currentTimeMillis(), new XPackLicenseStatus(operationMode, true, "Test license expired"));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@

package org.elasticsearch.xpack.esql.action;

import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.license.License;
import org.elasticsearch.license.XPackLicenseState;
import org.elasticsearch.license.internal.XPackLicenseStatus;
Expand All @@ -23,10 +22,6 @@
* - an expired enterprise or trial license
*/
public class EsqlPluginWithNonEnterpriseOrExpiredLicense extends EsqlPlugin {
public EsqlPluginWithNonEnterpriseOrExpiredLicense(Settings settings) {
super(settings);
}

protected XPackLicenseState getLicenseState() {
License.OperationMode operationMode;
boolean active;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@

package org.elasticsearch.xpack.esql.spatial;

import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.license.License;
import org.elasticsearch.license.XPackLicenseState;
import org.elasticsearch.license.internal.XPackLicenseStatus;
Expand Down Expand Up @@ -50,10 +49,6 @@ private static XPackLicenseState getLicenseState() {
* This is used to test the behavior of spatial functions when no valid license is present.
*/
public static class TestEsqlPlugin extends EsqlPlugin {
public TestEsqlPlugin(Settings settings) {
super(settings);
}

protected XPackLicenseState getLicenseState() {
return SpatialNoLicenseTestCase.getLicenseState();
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
/*
* 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.analysis;

import org.elasticsearch.cluster.project.ProjectResolver;
import org.elasticsearch.cluster.service.ClusterService;
import org.elasticsearch.xpack.esql.common.Failures;
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;

import java.util.List;
import java.util.function.BiConsumer;

/**
* SPI provider interface for supplying additional ESQL plan checks to be performed during verification.
*/
public interface PlanCheckerProvider {
/**
* Build a list of checks to perform on the plan. Each one is called once per
* {@link LogicalPlan} node in the plan.
*/
List<BiConsumer<LogicalPlan, Failures>> checkers(ProjectResolver projectResolver, ClusterService clusterService);
}
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@

package org.elasticsearch.xpack.esql.analysis;

import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.license.XPackLicenseState;
import org.elasticsearch.xpack.esql.LicenseAware;
import org.elasticsearch.xpack.esql.capabilities.PostAnalysisPlanVerificationAware;
Expand Down Expand Up @@ -56,31 +55,22 @@
* step does type resolution and fails queries based on invalid type expressions.
*/
public class Verifier {
public interface ExtraCheckers {
/**
* Build a list of checks to perform on the plan. Each one is called once per
* {@link LogicalPlan} node in the plan.
*/
List<BiConsumer<LogicalPlan, Failures>> extra(Settings settings);
}

/**
* Extra plan verification checks defined in plugins.
*/
private final List<ExtraCheckers> extraCheckers;
private final List<BiConsumer<LogicalPlan, Failures>> extraCheckers;
private final Metrics metrics;
private final XPackLicenseState licenseState;
private final Settings settings;

public Verifier(Metrics metrics, XPackLicenseState licenseState) {
this(metrics, licenseState, Collections.emptyList(), Settings.EMPTY);
this(metrics, licenseState, Collections.emptyList());
}

public Verifier(Metrics metrics, XPackLicenseState licenseState, List<ExtraCheckers> extraCheckers, Settings settings) {
public Verifier(Metrics metrics, XPackLicenseState licenseState, List<BiConsumer<LogicalPlan, Failures>> extraCheckers) {
this.metrics = metrics;
this.licenseState = licenseState;
this.extraCheckers = extraCheckers;
this.settings = settings;
}

/**
Expand All @@ -104,9 +94,7 @@ Collection<Failure> verify(LogicalPlan plan, BitSet partialMetrics) {

// collect plan checkers
var planCheckers = planCheckers(plan);
for (ExtraCheckers e : extraCheckers) {
planCheckers.addAll(e.extra(settings));
}
planCheckers.addAll(extraCheckers);

// Concrete verifications
plan.forEachDown(p -> {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,21 +8,22 @@
package org.elasticsearch.xpack.esql.execution;

import org.elasticsearch.action.ActionListener;
import org.elasticsearch.common.settings.Settings;
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;
import org.elasticsearch.xpack.esql.analysis.PreAnalyzer;
import org.elasticsearch.xpack.esql.analysis.Verifier;
import org.elasticsearch.xpack.esql.common.Failures;
import org.elasticsearch.xpack.esql.core.expression.FoldContext;
import org.elasticsearch.xpack.esql.enrich.EnrichPolicyResolver;
import org.elasticsearch.xpack.esql.expression.function.EsqlFunctionRegistry;
import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
import org.elasticsearch.xpack.esql.optimizer.LogicalPlanOptimizer;
import org.elasticsearch.xpack.esql.optimizer.LogicalPlanPreOptimizer;
import org.elasticsearch.xpack.esql.optimizer.LogicalPreOptimizerContext;
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
import org.elasticsearch.xpack.esql.planner.mapper.Mapper;
import org.elasticsearch.xpack.esql.plugin.TransportActionServices;
import org.elasticsearch.xpack.esql.querylog.EsqlQueryLog;
Expand All @@ -36,6 +37,7 @@
import org.elasticsearch.xpack.esql.telemetry.QueryMetric;

import java.util.List;
import java.util.function.BiConsumer;

import static org.elasticsearch.action.ActionListener.wrap;

Expand All @@ -55,15 +57,14 @@ public PlanExecutor(
MeterRegistry meterRegistry,
XPackLicenseState licenseState,
EsqlQueryLog queryLog,
List<Verifier.ExtraCheckers> extraCheckers,
Settings settings
List<BiConsumer<LogicalPlan, Failures>> extraCheckers
) {
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, licenseState, extraCheckers, settings);
this.verifier = new Verifier(metrics, licenseState, extraCheckers);
this.planTelemetryManager = new PlanTelemetryManager(meterRegistry);
this.queryLog = queryLog;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -67,14 +67,16 @@
import org.elasticsearch.xpack.esql.action.RestEsqlListQueriesAction;
import org.elasticsearch.xpack.esql.action.RestEsqlQueryAction;
import org.elasticsearch.xpack.esql.action.RestEsqlStopAsyncAction;
import org.elasticsearch.xpack.esql.analysis.Verifier;
import org.elasticsearch.xpack.esql.analysis.PlanCheckerProvider;
import org.elasticsearch.xpack.esql.common.Failures;
import org.elasticsearch.xpack.esql.enrich.EnrichLookupOperator;
import org.elasticsearch.xpack.esql.enrich.LookupFromIndexOperator;
import org.elasticsearch.xpack.esql.execution.PlanExecutor;
import org.elasticsearch.xpack.esql.expression.ExpressionWritables;
import org.elasticsearch.xpack.esql.io.stream.ExpressionQueryBuilder;
import org.elasticsearch.xpack.esql.io.stream.PlanStreamWrapperQueryBuilder;
import org.elasticsearch.xpack.esql.plan.PlanWritables;
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
import org.elasticsearch.xpack.esql.planner.PhysicalSettings;
import org.elasticsearch.xpack.esql.querydsl.query.SingleValueQuery;
import org.elasticsearch.xpack.esql.querylog.EsqlQueryLog;
Expand All @@ -85,6 +87,7 @@
import java.util.Collection;
import java.util.List;
import java.util.Objects;
import java.util.function.BiConsumer;
import java.util.function.Predicate;
import java.util.function.Supplier;

Expand Down Expand Up @@ -184,12 +187,7 @@ public class EsqlPlugin extends Plugin implements ActionPlugin, ExtensiblePlugin
Setting.Property.Dynamic
);

private final List<Verifier.ExtraCheckers> extraCheckers = new ArrayList<>();
private final Settings settings;

public EsqlPlugin(Settings settings) {
this.settings = settings;
}
private final List<PlanCheckerProvider> extraCheckerProviders = new ArrayList<>();

@Override
public Collection<?> createComponents(PluginServices services) {
Expand All @@ -203,14 +201,17 @@ public Collection<?> createComponents(PluginServices services) {
BigArrays bigArrays = services.indicesService().getBigArrays().withCircuitBreaking();
var blockFactoryProvider = blockFactoryProvider(circuitBreaker, bigArrays, maxPrimitiveArrayBlockSize);
setupSharedSecrets();
List<BiConsumer<LogicalPlan, Failures>> extraCheckers = extraCheckerProviders.stream()
.flatMap(p -> p.checkers(services.projectResolver(), services.clusterService()).stream())
.toList();

return List.of(
new PlanExecutor(
new IndexResolver(services.client()),
services.telemetryProvider().getMeterRegistry(),
getLicenseState(),
new EsqlQueryLog(services.clusterService().getClusterSettings(), services.slowLogFieldProvider()),
extraCheckers,
settings
extraCheckers
),
new ExchangeService(
services.clusterService().getSettings(),
Expand Down Expand Up @@ -349,6 +350,6 @@ public List<ExecutorBuilder<?>> getExecutorBuilders(Settings settings) {

@Override
public void loadExtensions(ExtensionLoader loader) {
extraCheckers.addAll(loader.loadExtensions(Verifier.ExtraCheckers.class));
extraCheckerProviders.addAll(loader.loadExtensions(PlanCheckerProvider.class));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ protected Writeable.Reader<ClusterComputeRequest> instanceReader() {
protected NamedWriteableRegistry getNamedWriteableRegistry() {
List<NamedWriteableRegistry.Entry> writeables = new ArrayList<>();
writeables.addAll(new SearchModule(Settings.EMPTY, List.of()).getNamedWriteables());
writeables.addAll(new EsqlPlugin(Settings.EMPTY).getNamedWriteables());
writeables.addAll(new EsqlPlugin().getNamedWriteables());
return new NamedWriteableRegistry(writeables);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ protected Writeable.Reader<DataNodeRequest> instanceReader() {
protected NamedWriteableRegistry getNamedWriteableRegistry() {
List<NamedWriteableRegistry.Entry> writeables = new ArrayList<>();
writeables.addAll(new SearchModule(Settings.EMPTY, List.of()).getNamedWriteables());
writeables.addAll(new EsqlPlugin(Settings.EMPTY).getNamedWriteables());
writeables.addAll(new EsqlPlugin().getNamedWriteables());
return new NamedWriteableRegistry(writeables);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -150,14 +150,7 @@ public void testFailedMetric() {
return null;
}).when(esqlClient).execute(eq(EsqlResolveFieldsAction.TYPE), any(), any());

var planExecutor = new PlanExecutor(
indexResolver,
MeterRegistry.NOOP,
new XPackLicenseState(() -> 0L),
mockQueryLog(),
List.of(),
Settings.EMPTY
);
var planExecutor = new PlanExecutor(indexResolver, MeterRegistry.NOOP, new XPackLicenseState(() -> 0L), mockQueryLog(), List.of());
var enrichResolver = mockEnrichResolver();

var request = new EsqlQueryRequest();
Expand Down