ESQL: Add extension point for extra plan checks

Author: nik9000

benchmarks/src/main/java/org/elasticsearch/benchmark/_nightly/esql/QueryPlanningBenchmark.java

@@ -47,6 +47,7 @@
 import org.openjdk.jmh.infra.Blackhole;
 
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
@@ -110,7 +111,7 @@ public void setup() {
                 new EnrichResolution(),
                 InferenceResolution.EMPTY
             ),
-            new Verifier(new Metrics(functionRegistry), new XPackLicenseState(() -> 0L))
+            new Verifier(List.of(), new Metrics(functionRegistry), new XPackLicenseState(() -> 0L))
         );
         defaultOptimizer = new LogicalPlanOptimizer(new LogicalOptimizerContext(config, FoldContext.small()));
     }

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

@@ -384,7 +384,7 @@ public static LogicalOptimizerContext unboundLogicalOptimizerContext() {
         return new LogicalOptimizerContext(EsqlTestUtils.TEST_CFG, FoldContext.small());
     }
 
-    public static final Verifier TEST_VERIFIER = new Verifier(new Metrics(new EsqlFunctionRegistry()), new XPackLicenseState(() -> 0L));
+    public static final Verifier TEST_VERIFIER = new Verifier(List.of(), new Metrics(new EsqlFunctionRegistry()), new XPackLicenseState(() -> 0L));
 
     public static final TransportActionServices MOCK_TRANSPORT_ACTION_SERVICES = new TransportActionServices(
         mock(TransportService.class),

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

@@ -54,11 +54,23 @@
  * 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();
+    }
 
+    /**
+     * Extra plan verification checks defined in plugins.
+     */
+    private final List<ExtraCheckers> extraCheckers;
     private final Metrics metrics;
     private final XPackLicenseState licenseState;
 
-    public Verifier(Metrics metrics, XPackLicenseState licenseState) {
+    public Verifier(List<ExtraCheckers> extraCheckers, Metrics metrics, XPackLicenseState licenseState) {
+        this.extraCheckers = extraCheckers;
         this.metrics = metrics;
         this.licenseState = licenseState;
     }
@@ -84,6 +96,9 @@ Collection<Failure> verify(LogicalPlan plan, BitSet partialMetrics) {
 
         // collect plan checkers
         var planCheckers = planCheckers(plan);
+        for (ExtraCheckers e : extraCheckers) {
+            planCheckers.addAll(e.extra());
+        }
 
         // Concrete verifications
         plan.forEachDown(p -> {
@@ -186,6 +201,9 @@ else if (p instanceof Lookup lookup) {
         });
     }
 
+    /**
+     * Build a list of checkers based on the components in the plan.
+     */
     private static List<BiConsumer<LogicalPlan, Failures>> planCheckers(LogicalPlan plan) {
         List<BiConsumer<LogicalPlan, Failures>> planCheckers = new ArrayList<>();
         Consumer<? super Node<?>> collectPlanCheckers = p -> {

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

@@ -15,11 +15,13 @@
 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.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;
@@ -32,6 +34,9 @@
 import org.elasticsearch.xpack.esql.telemetry.PlanTelemetryManager;
 import org.elasticsearch.xpack.esql.telemetry.QueryMetric;
 
+import java.util.List;
+import java.util.function.BiConsumer;
+
 import static org.elasticsearch.action.ActionListener.wrap;
 
 public class PlanExecutor {
@@ -44,16 +49,24 @@ public class PlanExecutor {
     private final Verifier verifier;
     private final PlanTelemetryManager planTelemetryManager;
     private final EsqlQueryLog queryLog;
+    private final List<Verifier.ExtraCheckers> extraCheckers;
 
-    public PlanExecutor(IndexResolver indexResolver, MeterRegistry meterRegistry, XPackLicenseState licenseState, EsqlQueryLog queryLog) {
+    public PlanExecutor(
+        IndexResolver indexResolver,
+        MeterRegistry meterRegistry,
+        XPackLicenseState licenseState,
+        EsqlQueryLog queryLog,
+        List<Verifier.ExtraCheckers> 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);
+        this.verifier = new Verifier(extraCheckers, metrics, licenseState);
         this.planTelemetryManager = new PlanTelemetryManager(meterRegistry);
         this.queryLog = queryLog;
+        this.extraCheckers = extraCheckers;
     }
 
     public void esql(

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

@@ -6,6 +6,7 @@
  */
 package org.elasticsearch.xpack.esql.plugin;
 
+import org.apache.lucene.util.SetOnce;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
 import org.elasticsearch.cluster.node.DiscoveryNodes;
 import org.elasticsearch.common.breaker.CircuitBreaker;
@@ -41,6 +42,7 @@
 import org.elasticsearch.features.NodeFeature;
 import org.elasticsearch.license.XPackLicenseState;
 import org.elasticsearch.plugins.ActionPlugin;
+import org.elasticsearch.plugins.ExtensiblePlugin;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.rest.RestController;
 import org.elasticsearch.rest.RestHandler;
@@ -66,6 +68,7 @@
 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.enrich.EnrichLookupOperator;
 import org.elasticsearch.xpack.esql.enrich.LookupFromIndexOperator;
 import org.elasticsearch.xpack.esql.execution.PlanExecutor;
@@ -83,7 +86,7 @@
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 
-public class EsqlPlugin extends Plugin implements ActionPlugin {
+public class EsqlPlugin extends Plugin implements ActionPlugin, ExtensiblePlugin {
     public static final boolean INLINESTATS_FEATURE_FLAG = new FeatureFlag("esql_inlinestats").isEnabled();
 
     public static final String ESQL_WORKER_THREAD_POOL_NAME = "esql_worker";
@@ -187,6 +190,8 @@ public class EsqlPlugin extends Plugin implements ActionPlugin {
         Setting.Property.Dynamic
     );
 
+    private SetOnce<List<Verifier.ExtraCheckers>> extraCheckers = new SetOnce<>();
+
     @Override
     public Collection<?> createComponents(PluginServices services) {
         CircuitBreaker circuitBreaker = services.indicesService().getBigArrays().breakerService().getBreaker("request");
@@ -204,7 +209,8 @@ public Collection<?> createComponents(PluginServices services) {
                 new IndexResolver(services.client()),
                 services.telemetryProvider().getMeterRegistry(),
                 getLicenseState(),
-                new EsqlQueryLog(services.clusterService().getClusterSettings(), services.slowLogFieldProvider())
+                new EsqlQueryLog(services.clusterService().getClusterSettings(), services.slowLogFieldProvider()),
+                extraCheckers.get()
             ),
             new ExchangeService(
                 services.clusterService().getSettings(),
@@ -334,4 +340,9 @@ public List<ExecutorBuilder<?>> getExecutorBuilders(Settings settings) {
             )
         );
     }
+
+    @Override
+    public void loadExtensions(ExtensionLoader loader) {
+        extraCheckers.set(loader.loadExtensions(Verifier.ExtraCheckers.class));
+    }
 }

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

@@ -98,7 +98,7 @@ private static Analyzer analyzer(EsqlFunctionRegistry registry, License.Operatio
                 defaultEnrichResolution(),
                 emptyInferenceResolution()
             ),
-            new Verifier(new Metrics(new EsqlFunctionRegistry()), getLicenseState(operationMode))
+            new Verifier(List.of(), new Metrics(new EsqlFunctionRegistry()), getLicenseState(operationMode))
         );
     }
 

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

@@ -235,7 +235,7 @@ private Analyzer makeAnalyzer(String mappingFileName, EnrichResolution enrichRes
                 enrichResolution,
                 emptyInferenceResolution()
             ),
-            new Verifier(new Metrics(new EsqlFunctionRegistry()), new XPackLicenseState(() -> 0L))
+            new Verifier(List.of(), new Metrics(new EsqlFunctionRegistry()), new XPackLicenseState(() -> 0L))
         );
     }
 

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

@@ -54,7 +54,7 @@ private static Analyzer makeAnalyzer(String mappingFileName) {
                 emptyPolicyResolution(),
                 emptyInferenceResolution()
             ),
-            new Verifier(new Metrics(new EsqlFunctionRegistry()), new XPackLicenseState(() -> 0L))
+            new Verifier(List.of(), new Metrics(new EsqlFunctionRegistry()), new XPackLicenseState(() -> 0L))
         );
     }
 

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

@@ -150,7 +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());
+        var planExecutor = new PlanExecutor(indexResolver, MeterRegistry.NOOP, new XPackLicenseState(() -> 0L), mockQueryLog(), List.of());
         var enrichResolver = mockEnrichResolver();
 
         var request = new EsqlQueryRequest();

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

@@ -206,7 +206,7 @@ public void testTwoWhereQuery() {
 
     public void testTwoQueriesExecuted() {
         Metrics metrics = new Metrics(new EsqlFunctionRegistry());
-        Verifier verifier = new Verifier(metrics, new XPackLicenseState(() -> 0L));
+        Verifier verifier = new Verifier(List.of(), metrics, new XPackLicenseState(() -> 0L));
         esqlWithVerifier("""
                from employees
                | where languages > 2
@@ -253,7 +253,7 @@ public void testTwoQueriesExecuted() {
 
     public void testMultipleFunctions() {
         Metrics metrics = new Metrics(new EsqlFunctionRegistry());
-        Verifier verifier = new Verifier(metrics, new XPackLicenseState(() -> 0L));
+        Verifier verifier = new Verifier(List.of(), metrics, new XPackLicenseState(() -> 0L));
         esqlWithVerifier("""
                from employees
                | where languages > 2
@@ -552,7 +552,7 @@ private Counters esql(String esql, Verifier v) {
         Metrics metrics = null;
         if (v == null) {
             metrics = new Metrics(new EsqlFunctionRegistry());
-            verifier = new Verifier(metrics, new XPackLicenseState(() -> 0L));
+            verifier = new Verifier(List.of(), metrics, new XPackLicenseState(() -> 0L));
         }
         analyzer(verifier).analyze(parser.createStatement(esql));