ES|QL: refactor query generator and move it to testFixtures

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();
@@ -104,9 +109,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;
@@ -133,25 +138,25 @@ 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, EsqlQueryGenerator.sourceCommand(), mappingInfo, exec);
+            EsqlQueryGenerator.generatePipeline(MAX_DEPTH, EsqlQueryGenerator.sourceCommand(), mappingInfo, exec, this);
         }
     }
 
     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,
@@ -172,7 +177,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;
@@ -192,35 +197,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 command, int depth) {
         try {
             Map<String, Object> json = RestEsqlTestCase.runEsql(
                 new RestEsqlTestCase.RequestObjectBuilder().query(command).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(command, depth, outputSchema, values, null);
         } catch (Exception e) {
-            return new EsqlQueryGenerator.QueryExecuted(command, depth, null, null, e);
+            return new QueryExecuted(command, 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(command, 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());
     }
 
@@ -240,10 +246,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,12 @@
+/*
+ * 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;
+
+public record Column(String name, String type, List<String> originalTypes) {}

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/EsqlQueryGenerator.java → x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/EsqlQueryGenerator.java

@@ -5,26 +5,26 @@
  * 2.0.
  */
 
-package org.elasticsearch.xpack.esql.qa.rest.generative;
+package org.elasticsearch.xpack.esql.generator;
 
 import org.elasticsearch.xpack.esql.CsvTestsDataLoader;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.ChangePointGenerator;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.DissectGenerator;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.DropGenerator;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.EnrichGenerator;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.EvalGenerator;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.ForkGenerator;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.GrokGenerator;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.KeepGenerator;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.LimitGenerator;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.LookupJoinGenerator;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.MvExpandGenerator;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.RenameGenerator;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.SortGenerator;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.StatsGenerator;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.WhereGenerator;
-import org.elasticsearch.xpack.esql.qa.rest.generative.command.source.FromGenerator;
+import org.elasticsearch.xpack.esql.generator.command.CommandGenerator;
+import org.elasticsearch.xpack.esql.generator.command.pipe.ChangePointGenerator;
+import org.elasticsearch.xpack.esql.generator.command.pipe.DissectGenerator;
+import org.elasticsearch.xpack.esql.generator.command.pipe.DropGenerator;
+import org.elasticsearch.xpack.esql.generator.command.pipe.EnrichGenerator;
+import org.elasticsearch.xpack.esql.generator.command.pipe.EvalGenerator;
+import org.elasticsearch.xpack.esql.generator.command.pipe.ForkGenerator;
+import org.elasticsearch.xpack.esql.generator.command.pipe.GrokGenerator;
+import org.elasticsearch.xpack.esql.generator.command.pipe.KeepGenerator;
+import org.elasticsearch.xpack.esql.generator.command.pipe.LimitGenerator;
+import org.elasticsearch.xpack.esql.generator.command.pipe.LookupJoinGenerator;
+import org.elasticsearch.xpack.esql.generator.command.pipe.MvExpandGenerator;
+import org.elasticsearch.xpack.esql.generator.command.pipe.RenameGenerator;
+import org.elasticsearch.xpack.esql.generator.command.pipe.SortGenerator;
+import org.elasticsearch.xpack.esql.generator.command.pipe.StatsGenerator;
+import org.elasticsearch.xpack.esql.generator.command.pipe.WhereGenerator;
+import org.elasticsearch.xpack.esql.generator.command.source.FromGenerator;
 
 import java.util.List;
 import java.util.Set;
@@ -42,10 +42,6 @@ public class EsqlQueryGenerator {
     public static final String COLUMN_TYPE = "type";
     public static final String COLUMN_ORIGINAL_TYPES = "original_types";
 
-    public record Column(String name, String type, List<String> originalTypes) {}
-
-    public record QueryExecuted(String query, int depth, List<Column> outputSchema, List<List<Object>> result, Exception exception) {}
-
     /**
      * These are commands that are at the beginning of the query, eg. FROM
      */
@@ -87,17 +83,18 @@ public interface Executor {
 
         boolean continueExecuting();
 
-        List<EsqlQueryGenerator.Column> currentSchema();
+        List<Column> currentSchema();
 
     }
 
     public static void generatePipeline(
         final int depth,
         CommandGenerator commandGenerator,
         final CommandGenerator.QuerySchema schema,
-        Executor executor
+        Executor executor,
+        QueryExecutor queryExecutor
     ) {
-        CommandGenerator.CommandDescription desc = commandGenerator.generate(List.of(), List.of(), schema);
+        CommandGenerator.CommandDescription desc = commandGenerator.generate(List.of(), List.of(), schema, queryExecutor);
         executor.run(commandGenerator, desc);
         if (executor.continueExecuting() == false) {
             return;
@@ -108,7 +105,7 @@ public static void generatePipeline(
                 break;
             }
             commandGenerator = EsqlQueryGenerator.randomPipeCommandGenerator();
-            desc = commandGenerator.generate(executor.previousCommands(), executor.currentSchema(), schema);
+            desc = commandGenerator.generate(executor.previousCommands(), executor.currentSchema(), schema, queryExecutor);
             if (desc == CommandGenerator.EMPTY_DESCRIPTION) {
                 continue;
             }
@@ -188,12 +185,12 @@ public static String randomGroupableName(List<Column> previousOutput) {
     }
 
     public static boolean groupable(Column col) {
-        return col.type.equals("keyword")
-            || col.type.equals("text")
-            || col.type.equals("long")
-            || col.type.equals("integer")
-            || col.type.equals("ip")
-            || col.type.equals("version");
+        return col.type().equals("keyword")
+            || col.type().equals("text")
+            || col.type().equals("long")
+            || col.type().equals("integer")
+            || col.type().equals("ip")
+            || col.type().equals("version");
     }
 
     /**
@@ -209,12 +206,12 @@ public static String randomSortableName(List<Column> previousOutput) {
     }
 
     public static boolean sortable(Column col) {
-        return col.type.equals("keyword")
-            || col.type.equals("text")
-            || col.type.equals("long")
-            || col.type.equals("integer")
-            || col.type.equals("ip")
-            || col.type.equals("version");
+        return col.type().equals("keyword")
+            || col.type().equals("text")
+            || col.type().equals("long")
+            || col.type().equals("integer")
+            || col.type().equals("ip")
+            || col.type().equals("version");
     }
 
     public static String agg(List<Column> previousOutput) {
@@ -308,7 +305,7 @@ public static String constantExpression() {
     /**
      * returns a random identifier or one of the existing names
      */
-    public static String randomAttributeOrIdentifier(List<EsqlQueryGenerator.Column> previousOutput) {
+    public static String randomAttributeOrIdentifier(List<Column> previousOutput) {
         String name;
         if (randomBoolean()) {
             name = EsqlQueryGenerator.randomIdentifier();
@@ -335,7 +332,7 @@ public static boolean fieldCanBeUsed(Column field) {
             || field.name().equals("<no-fields>")
             // no dense vectors for now, they are not supported in most commands
             || field.type().contains("vector")
-            || field.originalTypes.stream().anyMatch(x -> x.contains("vector"))) == false;
+            || field.originalTypes().stream().anyMatch(x -> x.contains("vector"))) == false;
     }
 
     public static String unquote(String colName) {

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

@@ -0,0 +1,12 @@
+/*
+ * 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;
+
+public record LookupIdx(String idxName, List<LookupIdxColumn> keys) {}

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

@@ -0,0 +1,10 @@
+/*
+ * 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;
+
+public record LookupIdxColumn(String name, String type) {}

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

@@ -0,0 +1,12 @@
+/*
+ * 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;
+
+public record QueryExecuted(String query, int depth, List<Column> outputSchema, List<List<Object>> result, Exception exception) {}

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

@@ -0,0 +1,12 @@
+/*
+ * 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;
+
+public interface QueryExecutor {
+    QueryExecuted execute(String command, int depth);
+}