ESQL: Add extension point for extra plan checks (#128840)

Author: nik9000

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/src/main/java/org/elasticsearch/xpack/esql/analysis/Verifier.java

@@ -39,6 +39,7 @@
 import java.util.ArrayList;
 import java.util.BitSet;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
@@ -54,13 +55,29 @@
  * 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) {
+        this(metrics, licenseState, Collections.emptyList());
+    }
+
+    public Verifier(Metrics metrics, XPackLicenseState licenseState, List<ExtraCheckers> extraCheckers) {
         this.metrics = metrics;
         this.licenseState = licenseState;
+        this.extraCheckers = extraCheckers;
     }
 
     /**
@@ -84,6 +101,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 +206,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(metrics, licenseState, extraCheckers);
         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(),
@@ -335,4 +340,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/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();