ES|QL: refactor query generator and move it to testFixtures (#134696)

Author: luigidellaquila

x-pack/plugin/esql/qa/server/single-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/GenerativeMetricsIT.java

@@ -11,9 +11,9 @@
 
 import org.elasticsearch.test.TestClustersThreadFilter;
 import org.elasticsearch.test.cluster.ElasticsearchCluster;
-import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator;
+import org.elasticsearch.xpack.esql.generator.EsqlQueryGenerator;
+import org.elasticsearch.xpack.esql.generator.command.CommandGenerator;
 import org.elasticsearch.xpack.esql.qa.rest.generative.GenerativeRestTest;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator;
 import org.junit.ClassRule;
 
 @ThreadLeakFilters(filters = TestClustersThreadFilter.class)

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java

@@ -9,13 +9,18 @@
 
 import org.elasticsearch.client.Request;
 import org.elasticsearch.client.ResponseException;
-import org.elasticsearch.core.Nullable;
 import org.elasticsearch.test.rest.ESRestTestCase;
 import org.elasticsearch.xpack.esql.AssertWarnings;
 import org.elasticsearch.xpack.esql.CsvTestsDataLoader;
+import org.elasticsearch.xpack.esql.generator.Column;
+import org.elasticsearch.xpack.esql.generator.EsqlQueryGenerator;
+import org.elasticsearch.xpack.esql.generator.LookupIdx;
+import org.elasticsearch.xpack.esql.generator.LookupIdxColumn;
+import org.elasticsearch.xpack.esql.generator.QueryExecuted;
+import org.elasticsearch.xpack.esql.generator.QueryExecutor;
+import org.elasticsearch.xpack.esql.generator.command.CommandGenerator;
 import org.elasticsearch.xpack.esql.qa.rest.ProfileLogger;
 import org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator;
 import org.junit.AfterClass;
 import org.junit.Before;
 import org.junit.Rule;
@@ -32,11 +37,11 @@
 import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.ENRICH_POLICIES;
 import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.availableDatasetsForEs;
 import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.loadDataSetIntoEs;
-import static org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator.COLUMN_NAME;
-import static org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator.COLUMN_ORIGINAL_TYPES;
-import static org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator.COLUMN_TYPE;
+import static org.elasticsearch.xpack.esql.generator.EsqlQueryGenerator.COLUMN_NAME;
+import static org.elasticsearch.xpack.esql.generator.EsqlQueryGenerator.COLUMN_ORIGINAL_TYPES;
+import static org.elasticsearch.xpack.esql.generator.EsqlQueryGenerator.COLUMN_TYPE;
 
-public abstract class GenerativeRestTest extends ESRestTestCase {
+public abstract class GenerativeRestTest extends ESRestTestCase implements QueryExecutor {
 
     @Rule(order = Integer.MIN_VALUE)
     public ProfileLogger profileLogger = new ProfileLogger();
@@ -56,7 +61,8 @@ public abstract class GenerativeRestTest extends ESRestTestCase {
 
         // Awaiting fixes for query failure
         "Unknown column \\[<all-fields-projected>\\]", // https://github.com/elastic/elasticsearch/issues/121741,
-        "Plan \\[ProjectExec\\[\\[<no-fields>.* optimized incorrectly due to missing references", // https://github.com/elastic/elasticsearch/issues/125866
+        // https://github.com/elastic/elasticsearch/issues/125866
+        "Plan \\[ProjectExec\\[\\[<no-fields>.* optimized incorrectly due to missing references",
         "The incoming YAML document exceeds the limit:", // still to investigate, but it seems to be specific to the test framework
         "Data too large", // Circuit breaker exceptions eg. https://github.com/elastic/elasticsearch/issues/130072
         "optimized incorrectly due to missing references", // https://github.com/elastic/elasticsearch/issues/131509
@@ -69,7 +75,8 @@ public abstract class GenerativeRestTest extends ESRestTestCase {
         "count_star .* can't be used with TS command",
         "time_series aggregate.* can only be used with the TS command",
         "Invalid call to dataType on an unresolved object \\?LASTOVERTIME", // https://github.com/elastic/elasticsearch/issues/134791
-        "class org.elasticsearch.compute.data..*Block cannot be cast to class org.elasticsearch.compute.data..*Block", // https://github.com/elastic/elasticsearch/issues/134793
+        // https://github.com/elastic/elasticsearch/issues/134793
+        "class org.elasticsearch.compute.data..*Block cannot be cast to class org.elasticsearch.compute.data..*Block",
         "Output has changed from \\[.*\\] to \\[.*\\]" // https://github.com/elastic/elasticsearch/issues/134794
     );
 
@@ -116,9 +123,9 @@ public void run(CommandGenerator generator, CommandGenerator.CommandDescription
                     previousCommands.add(current);
                     final String command = current.commandString();
 
-                    final EsqlQueryGenerator.QueryExecuted result = previousResult == null
-                        ? execute(command, 0, profileLogger)
-                        : execute(previousResult.query() + command, previousResult.depth(), profileLogger);
+                    final QueryExecuted result = previousResult == null
+                        ? execute(command, 0)
+                        : execute(previousResult.query() + command, previousResult.depth());
                     previousResult = result;
 
                     final boolean hasException = result.exception() != null;
@@ -145,16 +152,16 @@ public boolean continueExecuting() {
                 }
 
                 @Override
-                public List<EsqlQueryGenerator.Column> currentSchema() {
+                public List<Column> currentSchema() {
                     return currentSchema;
                 }
 
                 boolean continueExecuting;
-                List<EsqlQueryGenerator.Column> currentSchema;
+                List<Column> currentSchema;
                 final List<CommandGenerator.CommandDescription> previousCommands = new ArrayList<>();
-                EsqlQueryGenerator.QueryExecuted previousResult;
+                QueryExecuted previousResult;
             };
-            EsqlQueryGenerator.generatePipeline(MAX_DEPTH, sourceCommand(), mappingInfo, exec, requiresTimeSeries());
+            EsqlQueryGenerator.generatePipeline(MAX_DEPTH, sourceCommand(), mappingInfo, exec, requiresTimeSeries(), this);
         }
     }
 
@@ -166,8 +173,8 @@ private static CommandGenerator.ValidationResult checkResults(
         List<CommandGenerator.CommandDescription> previousCommands,
         CommandGenerator commandGenerator,
         CommandGenerator.CommandDescription commandDescription,
-        EsqlQueryGenerator.QueryExecuted previousResult,
-        EsqlQueryGenerator.QueryExecuted result
+        QueryExecuted previousResult,
+        QueryExecuted result
     ) {
         CommandGenerator.ValidationResult outputValidation = commandGenerator.validateOutput(
             previousCommands,
@@ -188,7 +195,7 @@ private static CommandGenerator.ValidationResult checkResults(
         return outputValidation;
     }
 
-    private void checkException(EsqlQueryGenerator.QueryExecuted query) {
+    private void checkException(QueryExecuted query) {
         for (Pattern allowedError : ALLOWED_ERROR_PATTERNS) {
             if (isAllowedError(query.exception().getMessage(), allowedError)) {
                 return;
@@ -208,35 +215,36 @@ private static boolean isAllowedError(String errorMessage, Pattern allowedPatter
         return allowedPattern.matcher(errorWithoutLineBreaks).matches();
     }
 
+    @Override
     @SuppressWarnings("unchecked")
-    public static EsqlQueryGenerator.QueryExecuted execute(String command, int depth, @Nullable ProfileLogger profileLogger) {
+    public QueryExecuted execute(String query, int depth) {
         try {
             Map<String, Object> json = RestEsqlTestCase.runEsql(
-                new RestEsqlTestCase.RequestObjectBuilder().query(command).build(),
+                new RestEsqlTestCase.RequestObjectBuilder().query(query).build(),
                 new AssertWarnings.AllowedRegexes(List.of(Pattern.compile(".*"))),// we don't care about warnings
                 profileLogger,
                 RestEsqlTestCase.Mode.SYNC
             );
-            List<EsqlQueryGenerator.Column> outputSchema = outputSchema(json);
+            List<Column> outputSchema = outputSchema(json);
             List<List<Object>> values = (List<List<Object>>) json.get("values");
-            return new EsqlQueryGenerator.QueryExecuted(command, depth, outputSchema, values, null);
+            return new QueryExecuted(query, depth, outputSchema, values, null);
         } catch (Exception e) {
-            return new EsqlQueryGenerator.QueryExecuted(command, depth, null, null, e);
+            return new QueryExecuted(query, depth, null, null, e);
         } catch (AssertionError ae) {
             // this is for ensureNoWarnings()
-            return new EsqlQueryGenerator.QueryExecuted(command, depth, null, null, new RuntimeException(ae.getMessage()));
+            return new QueryExecuted(query, depth, null, null, new RuntimeException(ae.getMessage()));
         }
 
     }
 
     @SuppressWarnings("unchecked")
-    private static List<EsqlQueryGenerator.Column> outputSchema(Map<String, Object> a) {
+    private static List<Column> outputSchema(Map<String, Object> a) {
         List<Map<String, ?>> cols = (List<Map<String, ?>>) a.get("columns");
         if (cols == null) {
             return null;
         }
         return cols.stream()
-            .map(x -> new EsqlQueryGenerator.Column((String) x.get(COLUMN_NAME), (String) x.get(COLUMN_TYPE), originalTypes(x)))
+            .map(x -> new Column((String) x.get(COLUMN_NAME), (String) x.get(COLUMN_TYPE), originalTypes(x)))
             .collect(Collectors.toList());
     }
 
@@ -256,10 +264,6 @@ private List<String> availableIndices() throws IOException {
             .toList();
     }
 
-    public record LookupIdxColumn(String name, String type) {}
-
-    public record LookupIdx(String idxName, List<LookupIdxColumn> keys) {}
-
     private List<LookupIdx> lookupIndices() {
         List<LookupIdx> result = new ArrayList<>();
         // we don't have key info from the dataset loader, let's hardcode it for now

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/README.asciidoc

@@ -26,18 +26,11 @@ The result check happens at two levels:
 
 == Implementing your own command generator
 
-If you implement a new command, and you want it to be tested by the generative tests, you can add a command generator here.
+If you create a new command, we strongly recommend you to implement a command generator for it, so that it is tested
+by the generative tests.
 
-All you have to do is:
+To do so, see the documentation in the
+`x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/README.asciidoc` file.
 
-* add a class in `org.elasticsearch.xpack.esql.qa.rest.generative.command.source` (if it's a source command) or in `org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe` (if it's a pipe command)
-* Implement `CommandGenerator` interface (see its javadoc, it should be explicative. Or just have a look at one of the existing commands, eg. `SortGenerator`)
-** Implement `CommandGenerator.generate()` method, that will return the command.
-*** Have a look at `EsqlQueryGenerator`, it contains many utility methods that will help you generate random expressions.
-** Implement `CommandGenerator.validateOutput()` to validate the output of the query.
-* Add your class to `EsqlQueryGenerator.SOURCE_COMMANDS` (if it's a source command) or `EsqlQueryGenerator.PIPE_COMMANDS` (if it's a pipe command).
-* Run `GenerativeIT` at least a couple of times: these tests can be pretty noisy.
-* If you get unexpected errors (real bugs in ES|QL), please open an issue and add the error to `GenerativeRestTest.ALLOWED_ERRORS`. Run tests again until everything works fine.
 
 
-IMPORTANT: be careful when validating the output (Eg. the row count), as ES|QL can be quite non-deterministic when there are no SORTs

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

@@ -0,0 +1,18 @@
+/*
+ * 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.generator;
+
+import java.util.List;
+
+/**
+ * A column in the output schema of a query execution.
+ * @param name the field name
+ * @param type the field type
+ * @param originalTypes the original types, in case of a union type.
+ */
+public record Column(String name, String type, List<String> originalTypes) {}