ESQL: Add extension point for extra plan checks

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/server/extra-checkers/build.gradle

@@ -0,0 +1,29 @@
+apply plugin: 'elasticsearch.internal-es-plugin'
+apply plugin: 'elasticsearch.internal-java-rest-test'
+apply plugin: org.elasticsearch.gradle.internal.precommit.CheckstylePrecommitPlugin
+apply plugin: org.elasticsearch.gradle.internal.precommit.ForbiddenApisPrecommitPlugin
+apply plugin: org.elasticsearch.gradle.internal.precommit.ForbiddenPatternsPrecommitPlugin
+apply plugin: org.elasticsearch.gradle.internal.precommit.FilePermissionsPrecommitPlugin
+apply plugin: org.elasticsearch.gradle.internal.precommit.LoggerUsagePrecommitPlugin
+apply plugin: org.elasticsearch.gradle.internal.precommit.TestingConventionsPrecommitPlugin
+
+
+esplugin {
+  name = 'extra-checkers'
+  description = 'An example plugin disallowing CATEGORIZE'
+  classname ='org.elasticsearch.xpack.esql.qa.extra.ExtraCheckersPlugin'
+  extendedPlugins = ['x-pack-esql']
+}
+
+dependencies {
+  compileOnly project(':x-pack:plugin:esql')
+  compileOnly project(':x-pack:plugin:esql-core')
+  clusterPlugins project(':x-pack:plugin:esql:qa:server:extra-checkers')
+}
+
+tasks.named('javaRestTest') {
+  usesDefaultDistribution("to be triaged")
+  maxParallelForks = 1
+  jvmArgs('--add-opens=java.base/java.nio=ALL-UNNAMED')
+}
+

x-pack/plugin/esql/qa/server/extra-checkers/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/extra/ExtraCheckersIT.java

@@ -0,0 +1,80 @@
+/*
+ * 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.qa.extra;
+
+import com.carrotsearch.randomizedtesting.annotations.ThreadLeakFilters;
+
+import org.apache.http.util.EntityUtils;
+import org.elasticsearch.client.Request;
+import org.elasticsearch.client.Response;
+import org.elasticsearch.client.ResponseException;
+import org.elasticsearch.test.TestClustersThreadFilter;
+import org.elasticsearch.test.cluster.ElasticsearchCluster;
+import org.elasticsearch.test.cluster.local.distribution.DistributionType;
+import org.elasticsearch.test.rest.ESRestTestCase;
+import org.junit.ClassRule;
+
+import java.io.IOException;
+
+import static org.hamcrest.Matchers.containsString;
+import static org.hamcrest.Matchers.equalTo;
+
+@ThreadLeakFilters(filters = TestClustersThreadFilter.class)
+public class ExtraCheckersIT extends ESRestTestCase {
+    @ClassRule
+    public static ElasticsearchCluster cluster = ElasticsearchCluster.local()
+        .distribution(DistributionType.DEFAULT)
+        .setting("xpack.security.enabled", "false")
+        .setting("xpack.license.self_generated.type", "trial")
+        .shared(true)
+        .plugin("extra-checkers")
+        .build();
+
+    public void testWithCategorize() {
+        ResponseException e = expectThrows(ResponseException.class, () -> runEsql("""
+            {
+              "query": "ROW message=\\"foo bar\\" | STATS COUNT(*) BY CATEGORIZE(message) | LIMIT 1"
+            }"""));
+        assertThat(e.getResponse().getStatusLine().getStatusCode(), equalTo(400));
+        assertThat(e.getMessage(), containsString("line 1:43: CATEGORIZE is unsupported"));
+    }
+
+    public void testWithoutCategorize() throws IOException {
+        String result = runEsql("""
+            {
+              "query": "ROW message=\\"foo bar\\" | STATS COUNT(*) | LIMIT 1"
+            }""");
+        assertThat(result, containsString("""
+              "columns" : [
+                {
+                  "name" : "COUNT(*)",
+                  "type" : "long"
+                }
+              ],
+              "values" : [
+                [
+                  1
+                ]
+              ]
+            """));
+    }
+
+    private String runEsql(String json) throws IOException {
+        Request request = new Request("POST", "/_query");
+        request.setJsonEntity(json);
+        request.addParameter("error_trace", "");
+        request.addParameter("pretty", "");
+        Response response = client().performRequest(request);
+        return EntityUtils.toString(response.getEntity());
+    }
+
+    @Override
+    protected String getTestRestCluster() {
+        return cluster.getHttpAddresses();
+    }
+}

x-pack/plugin/esql/qa/server/extra-checkers/src/main/java/org/elasticsearch/xpack/esql/qa/extra/DisallowCategorize.java

@@ -0,0 +1,31 @@
+/*
+ * 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.qa.extra;
+
+import org.elasticsearch.xpack.esql.analysis.Verifier;
+import org.elasticsearch.xpack.esql.common.Failure;
+import org.elasticsearch.xpack.esql.common.Failures;
+import org.elasticsearch.xpack.esql.expression.function.grouping.Categorize;
+import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
+import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+
+import java.util.List;
+import java.util.function.BiConsumer;
+
+public class DisallowCategorize implements Verifier.ExtraCheckers {
+    @Override
+    public List<BiConsumer<LogicalPlan, Failures>> extra() {
+        return List.of(DisallowCategorize::disallowCategorize);
+    }
+
+    private static void disallowCategorize(LogicalPlan plan, Failures failures) {
+        if (plan instanceof Aggregate) {
+            plan.forEachExpression(Categorize.class, cat -> failures.add(new Failure(cat, "CATEGORIZE is unsupported")));
+        }
+    }
+}

x-pack/plugin/esql/qa/server/extra-checkers/src/main/java/org/elasticsearch/xpack/esql/qa/extra/ExtraCheckersPlugin.java

@@ -0,0 +1,16 @@
+/*
+ * 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.qa.extra;
+
+import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.xpack.esql.analysis.Verifier;
+
+/**
+ * Marker plugin to enable {@link Verifier.ExtraCheckers}.
+ */
+public class ExtraCheckersPlugin extends Plugin {}

x-pack/plugin/esql/qa/server/extra-checkers/src/main/resources/META-INF/services/org.elasticsearch.xpack.esql.analysis.Verifier$ExtraCheckers

@@ -0,0 +1,8 @@
+#
+# 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.
+#
+
+org.elasticsearch.xpack.esql.qa.extra.DisallowCategorize

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

@@ -384,7 +384,11 @@ 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

@@ -32,6 +32,8 @@
 import org.elasticsearch.xpack.esql.telemetry.PlanTelemetryManager;
 import org.elasticsearch.xpack.esql.telemetry.QueryMetric;
 
+import java.util.List;
+
 import static org.elasticsearch.action.ActionListener.wrap;
 
 public class PlanExecutor {
@@ -45,13 +47,19 @@ public class PlanExecutor {
     private final PlanTelemetryManager planTelemetryManager;
     private final EsqlQueryLog queryLog;
 
-    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;
     }

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

@@ -41,6 +41,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 +67,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 +85,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 +189,8 @@ public class EsqlPlugin extends Plugin implements ActionPlugin {
         Setting.Property.Dynamic
     );
 
+    private final List<Verifier.ExtraCheckers> extraCheckers = new ArrayList<>();
+
     @Override
     public Collection<?> createComponents(PluginServices services) {
         CircuitBreaker circuitBreaker = services.indicesService().getBigArrays().breakerService().getBreaker("request");
@@ -204,7 +208,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
             ),
             new ExchangeService(
                 services.clusterService().getSettings(),
@@ -334,4 +339,9 @@ public List<ExecutorBuilder<?>> getExecutorBuilders(Settings settings) {
             )
         );
     }
+
+    @Override
+    public void loadExtensions(ExtensionLoader loader) {
+        extraCheckers.addAll(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));