diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/EsqlQueryGenerator.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/EsqlQueryGenerator.java index 31fddae7c6859..46f48963c820c 100644 --- a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/EsqlQueryGenerator.java +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/EsqlQueryGenerator.java @@ -8,12 +8,23 @@ package org.elasticsearch.xpack.esql.qa.rest.generative; 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.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.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 java.util.ArrayList; -import java.util.HashMap; -import java.util.HashSet; import java.util.List; -import java.util.Map; import java.util.Set; import java.util.stream.Collectors; @@ -27,85 +38,41 @@ public class EsqlQueryGenerator { public record Column(String name, String type) {} - public record QueryExecuted(String query, int depth, List outputSchema, Exception exception) {} - - public static String sourceCommand(List availabeIndices) { - return switch (randomIntBetween(0, 1)) { - case 0 -> from(availabeIndices); - // case 1 -> metaFunctions(); - default -> from(availabeIndices); - // TODO re-enable ROW. - // now it crashes nodes in some cases: exiting java.lang.AssertionError: estimated row size [0] wasn't set - // default -> row(); - }; - - } + public record QueryExecuted(String query, int depth, List outputSchema, List> result, Exception exception) {} /** - * @param previousOutput a list of fieldName+type - * @param policies - * @return a new command that can process it as input + * These are commands that are at the beginning of the query, eg. FROM */ - public static String pipeCommand( - List previousOutput, - List policies, - List lookupIndices - ) { - return switch (randomIntBetween(0, 12)) { - case 0 -> dissect(previousOutput); - case 1 -> drop(previousOutput); - case 2 -> enrich(previousOutput, policies); - case 3 -> eval(previousOutput); - case 4 -> grok(previousOutput); - case 5 -> keep(previousOutput); - case 6 -> limit(); - case 7 -> mvExpand(previousOutput); - case 8 -> rename(previousOutput); - case 9 -> sort(previousOutput); - case 10 -> stats(previousOutput); - case 11 -> join(previousOutput, lookupIndices); - default -> where(previousOutput); - }; - } - - private static String join(List previousOutput, List lookupIndices) { - - GenerativeRestTest.LookupIdx lookupIdx = randomFrom(lookupIndices); - String lookupIdxName = lookupIdx.idxName(); - String idxKey = lookupIdx.key(); - String keyType = lookupIdx.keyType(); - - var candidateKeys = previousOutput.stream().filter(x -> x.type.equals(keyType)).toList(); - if (candidateKeys.isEmpty()) { - return ""; - } - Column key = randomFrom(candidateKeys); - return "| rename " + key.name + " as " + idxKey + " | lookup join " + lookupIdxName + " on " + idxKey; - } - - private static String where(List previousOutput) { - // TODO more complex conditions - StringBuilder result = new StringBuilder(" | where "); - int nConditions = randomIntBetween(1, 5); - for (int i = 0; i < nConditions; i++) { - String exp = booleanExpression(previousOutput); - if (exp == null) { - // cannot generate expressions, just skip - return ""; - } - if (i > 0) { - result.append(randomBoolean() ? " AND " : " OR "); - } - if (randomBoolean()) { - result.append(" NOT "); - } - result.append(exp); - } - - return result.toString(); - } + static List SOURCE_COMMANDS = List.of(FromGenerator.INSTANCE); - private static String booleanExpression(List previousOutput) { + /** + * These are downstream commands, ie. that cannot appear as the first command in a query + */ + static List PIPE_COMMANDS = List.of( + DissectGenerator.INSTANCE, + DropGenerator.INSTANCE, + EnrichGenerator.INSTANCE, + EvalGenerator.INSTANCE, + GrokGenerator.INSTANCE, + KeepGenerator.INSTANCE, + LimitGenerator.INSTANCE, + LookupJoinGenerator.INSTANCE, + MvExpandGenerator.INSTANCE, + RenameGenerator.INSTANCE, + SortGenerator.INSTANCE, + StatsGenerator.INSTANCE, + WhereGenerator.INSTANCE + ); + + public static CommandGenerator sourceCommand() { + return randomFrom(SOURCE_COMMANDS); + } + + public static CommandGenerator randomPipeCommandGenerator() { + return randomFrom(PIPE_COMMANDS); + } + + public static String booleanExpression(List previousOutput) { // TODO LIKE, RLIKE, functions etc. return switch (randomIntBetween(0, 3)) { case 0 -> { @@ -120,7 +87,7 @@ private static String booleanExpression(List previousOutput) { }; } - private static String mathCompareOperator() { + public static String mathCompareOperator() { return switch (randomIntBetween(0, 5)) { case 0 -> "=="; case 1 -> ">"; @@ -131,105 +98,12 @@ private static String mathCompareOperator() { }; } - private static String enrich(List previousOutput, List policies) { - String field = randomKeywordField(previousOutput); - if (field == null || policies.isEmpty()) { - return ""; - } - - // TODO add WITH - return " | enrich " + randomFrom(policiesOnKeyword(policies)).policyName() + " on " + field; - } - - private static List policiesOnKeyword(List policies) { + public static List policiesOnKeyword(List policies) { // TODO make it smarter and extend it to other types return policies.stream().filter(x -> Set.of("languages_policy").contains(x.policyName())).toList(); } - private static String grok(List previousOutput) { - String field = randomStringField(previousOutput); - if (field == null) { - return "";// no strings to grok, just skip - } - StringBuilder result = new StringBuilder(" | grok "); - result.append(field); - result.append(" \""); - for (int i = 0; i < randomIntBetween(1, 3); i++) { - if (i > 0) { - result.append(" "); - } - result.append("%{WORD:"); - if (randomBoolean()) { - result.append(randomIdentifier()); - } else { - String fieldName = randomRawName(previousOutput); - if (fieldName == null) { - fieldName = randomIdentifier(); - } - result.append(fieldName); - } - result.append("}"); - } - result.append("\""); - return result.toString(); - } - - private static String dissect(List previousOutput) { - String field = randomStringField(previousOutput); - if (field == null) { - return "";// no strings to dissect, just skip - } - StringBuilder result = new StringBuilder(" | dissect "); - result.append(field); - result.append(" \""); - for (int i = 0; i < randomIntBetween(1, 3); i++) { - if (i > 0) { - result.append(" "); - } - result.append("%{"); - if (randomBoolean()) { - result.append(randomIdentifier()); - } else { - String fieldName = randomRawName(previousOutput); - if (fieldName == null) { - fieldName = randomIdentifier(); - } - result.append(fieldName); - } - result.append("}"); - } - result.append("\""); - return result.toString(); - } - - private static String keep(List previousOutput) { - int n = randomIntBetween(1, previousOutput.size()); - Set proj = new HashSet<>(); - for (int i = 0; i < n; i++) { - if (randomIntBetween(0, 100) < 5) { - proj.add("*"); - } else { - String name = randomName(previousOutput); - if (name == null) { - continue; - } - if (name.length() > 1 && name.startsWith("`") == false && randomIntBetween(0, 100) < 10) { - if (randomBoolean()) { - name = name.substring(0, randomIntBetween(1, name.length() - 1)) + "*"; - } else { - name = "*" + name.substring(randomIntBetween(1, name.length() - 1)); - } - } - proj.add(name); - } - } - if (proj.isEmpty()) { - return ""; - } - return " | keep " + proj.stream().collect(Collectors.joining(", ")); - } - - private static String randomName(List previousOutput) { + public static String randomName(List previousOutput) { String result = randomRawName(previousOutput); if (result == null) { return null; @@ -244,7 +118,7 @@ private static String randomName(List previousOutput) { * Returns a field name from a list of columns. * Could be null if none of the fields can be considered */ - private static String randomRawName(List previousOutput) { + public static String randomRawName(List previousOutput) { var list = previousOutput.stream().filter(EsqlQueryGenerator::fieldCanBeUsed).toList(); if (list.isEmpty()) { return null; @@ -257,7 +131,7 @@ private static String randomRawName(List previousOutput) { * Returns a field that can be used for grouping. * Can return null */ - private static String randomGroupableName(List previousOutput) { + public static String randomGroupableName(List previousOutput) { var candidates = previousOutput.stream().filter(EsqlQueryGenerator::groupable).filter(EsqlQueryGenerator::fieldCanBeUsed).toList(); if (candidates.isEmpty()) { return null; @@ -265,7 +139,7 @@ private static String randomGroupableName(List previousOutput) { return randomFrom(candidates).name(); } - private static boolean groupable(Column col) { + public static boolean groupable(Column col) { return col.type.equals("keyword") || col.type.equals("text") || col.type.equals("long") @@ -278,7 +152,7 @@ private static boolean groupable(Column col) { * returns a field that can be sorted. * Null if no fields are sortable. */ - private static String randomSortableName(List previousOutput) { + public static String randomSortableName(List previousOutput) { var candidates = previousOutput.stream().filter(EsqlQueryGenerator::sortable).filter(EsqlQueryGenerator::fieldCanBeUsed).toList(); if (candidates.isEmpty()) { return null; @@ -286,7 +160,7 @@ private static String randomSortableName(List previousOutput) { return randomFrom(candidates).name(); } - private static boolean sortable(Column col) { + public static boolean sortable(Column col) { return col.type.equals("keyword") || col.type.equals("text") || col.type.equals("long") @@ -295,170 +169,7 @@ private static boolean sortable(Column col) { || col.type.equals("version"); } - private static String rename(List previousOutput) { - int n = randomIntBetween(1, Math.min(3, previousOutput.size())); - List proj = new ArrayList<>(); - - Map nameToType = new HashMap<>(); - for (Column column : previousOutput) { - nameToType.put(column.name, column.type); - } - List names = new ArrayList<>( - previousOutput.stream().filter(EsqlQueryGenerator::fieldCanBeUsed).map(Column::name).collect(Collectors.toList()) - ); - if (names.isEmpty()) { - return ""; - } - for (int i = 0; i < n; i++) { - if (names.isEmpty()) { - break; - } - var name = randomFrom(names); - if (nameToType.get(name).endsWith("_range")) { - // ranges are not fully supported yet - continue; - } - names.remove(name); - - String newName; - if (names.isEmpty() || randomBoolean()) { - newName = randomIdentifier(); - names.add(newName); - } else { - newName = names.get(randomIntBetween(0, names.size() - 1)); - } - nameToType.put(newName, nameToType.get(name)); - if (randomBoolean() && name.startsWith("`") == false) { - name = "`" + name + "`"; - } - if (randomBoolean() && newName.startsWith("`") == false) { - newName = "`" + newName + "`"; - } - proj.add(name + " AS " + newName); - } - if (proj.isEmpty()) { - return ""; - } - return " | rename " + proj.stream().collect(Collectors.joining(", ")); - } - - private static String drop(List previousOutput) { - if (previousOutput.size() < 2) { - return ""; // don't drop all of them, just do nothing - } - int n = randomIntBetween(1, previousOutput.size() - 1); - Set proj = new HashSet<>(); - for (int i = 0; i < n; i++) { - String name = randomRawName(previousOutput); - if (name == null) { - continue; - } - if (name.length() > 1 && name.startsWith("`") == false && randomIntBetween(0, 100) < 10) { - if (randomBoolean()) { - name = name.substring(0, randomIntBetween(1, name.length() - 1)) + "*"; - } else { - name = "*" + name.substring(randomIntBetween(1, name.length() - 1)); - } - } else if (name.startsWith("`") == false && (randomBoolean() || name.isEmpty())) { - name = "`" + name + "`"; - } - proj.add(name); - } - if (proj.isEmpty()) { - return ""; - } - return " | drop " + proj.stream().collect(Collectors.joining(", ")); - } - - private static String sort(List previousOutput) { - int n = randomIntBetween(1, previousOutput.size()); - Set proj = new HashSet<>(); - for (int i = 0; i < n; i++) { - String col = randomSortableName(previousOutput); - if (col == null) { - return "";// no sortable columns - } - proj.add(col); - } - return " | sort " - + proj.stream() - .map(x -> x + randomFrom("", " ASC", " DESC") + randomFrom("", " NULLS FIRST", " NULLS LAST")) - .collect(Collectors.joining(", ")); - } - - private static String mvExpand(List previousOutput) { - String toExpand = randomName(previousOutput); - if (toExpand == null) { - return ""; // no columns to expand - } - return " | mv_expand " + toExpand; - } - - private static String eval(List previousOutput) { - StringBuilder cmd = new StringBuilder(" | eval "); - int nFields = randomIntBetween(1, 10); - // TODO pass newly created fields to next expressions - for (int i = 0; i < nFields; i++) { - String name; - if (randomBoolean()) { - name = randomIdentifier(); - } else { - name = randomName(previousOutput); - if (name == null) { - name = randomIdentifier(); - } - } - String expression = expression(previousOutput); - if (i > 0) { - cmd.append(","); - } - cmd.append(" "); - cmd.append(name); - cmd.append(" = "); - cmd.append(expression); - } - return cmd.toString(); - } - - private static String stats(List previousOutput) { - List nonNull = previousOutput.stream() - .filter(EsqlQueryGenerator::fieldCanBeUsed) - .filter(x -> x.type().equals("null") == false) - .collect(Collectors.toList()); - if (nonNull.isEmpty()) { - return ""; // cannot do any stats, just skip - } - StringBuilder cmd = new StringBuilder(" | stats "); - int nStats = randomIntBetween(1, 5); - for (int i = 0; i < nStats; i++) { - String name; - if (randomBoolean()) { - name = randomIdentifier(); - } else { - name = randomName(previousOutput); - if (name == null) { - name = randomIdentifier(); - } - } - String expression = agg(nonNull); - if (i > 0) { - cmd.append(","); - } - cmd.append(" "); - cmd.append(name); - cmd.append(" = "); - cmd.append(expression); - } - if (randomBoolean()) { - var col = randomGroupableName(nonNull); - if (col != null) { - cmd.append(" by " + col); - } - } - return cmd.toString(); - } - - private static String agg(List previousOutput) { + public static String agg(List previousOutput) { String name = randomNumericOrDateField(previousOutput); if (name != null && randomBoolean()) { // numerics only @@ -480,23 +191,23 @@ private static String agg(List previousOutput) { }; } - private static String randomNumericOrDateField(List previousOutput) { + public static String randomNumericOrDateField(List previousOutput) { return randomName(previousOutput, Set.of("long", "integer", "double", "date")); } - private static String randomNumericField(List previousOutput) { + public static String randomNumericField(List previousOutput) { return randomName(previousOutput, Set.of("long", "integer", "double")); } - private static String randomStringField(List previousOutput) { + public static String randomStringField(List previousOutput) { return randomName(previousOutput, Set.of("text", "keyword")); } - private static String randomKeywordField(List previousOutput) { + public static String randomKeywordField(List previousOutput) { return randomName(previousOutput, Set.of("keyword")); } - private static String randomName(List cols, Set allowedTypes) { + public static String randomName(List cols, Set allowedTypes) { List items = cols.stream().filter(x -> allowedTypes.contains(x.type())).map(Column::name).collect(Collectors.toList()); if (items.size() == 0) { return null; @@ -504,37 +215,16 @@ private static String randomName(List cols, Set allowedTypes) { return items.get(randomIntBetween(0, items.size() - 1)); } - private static String expression(List previousOutput) { + public static String expression(List previousOutput) { // TODO improve!!! return constantExpression(); } - public static String limit() { - return " | limit " + randomIntBetween(0, 15000); - } - - private static String from(List availabeIndices) { - StringBuilder result = new StringBuilder("from "); - int items = randomIntBetween(1, 3); - for (int i = 0; i < items; i++) { - String pattern = indexPattern(availabeIndices.get(randomIntBetween(0, availabeIndices.size() - 1))); - if (i > 0) { - result.append(","); - } - result.append(pattern); - } - return result.toString(); - } - - private static String metaFunctions() { - return "meta functions"; - } - - private static String indexPattern(String indexName) { + public static String indexPattern(String indexName) { return randomBoolean() ? indexName : indexName.substring(0, randomIntBetween(0, indexName.length())) + "*"; } - private static String row() { + public static String row() { StringBuilder cmd = new StringBuilder("row "); int nFields = randomIntBetween(1, 10); for (int i = 0; i < nFields; i++) { @@ -551,7 +241,7 @@ private static String row() { return cmd.toString(); } - private static String constantExpression() { + public static String constantExpression() { // TODO not only simple values, but also foldable expressions return switch (randomIntBetween(0, 4)) { case 0 -> "" + randomIntBetween(Integer.MIN_VALUE, Integer.MAX_VALUE); @@ -563,13 +253,13 @@ private static String constantExpression() { } - private static String randomIdentifier() { + public static String randomIdentifier() { // Let's create identifiers that are long enough to avoid collisions with reserved keywords. // There could be a smarter way (introspection on the lexer class?), but probably it's not worth the effort return randomAlphaOfLength(randomIntBetween(8, 12)); } - private static boolean fieldCanBeUsed(Column field) { + public static boolean fieldCanBeUsed(Column field) { return ( // https://github.com/elastic/elasticsearch/issues/121741 field.name().equals("") @@ -577,4 +267,11 @@ private static boolean fieldCanBeUsed(Column field) { || field.name().equals("")) == false; } + public static String unquote(String colName) { + if (colName.startsWith("`") && colName.endsWith("`")) { + return colName.substring(1, colName.length() - 1); + } + return colName; + } + } diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java index 8ae28477e03bb..88df38cc347cf 100644 --- a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java @@ -12,6 +12,7 @@ import org.elasticsearch.test.rest.ESRestTestCase; import org.elasticsearch.xpack.esql.CsvTestsDataLoader; 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; @@ -47,11 +48,14 @@ public abstract class GenerativeRestTest extends ESRestTestCase { "Field '.*' shadowed by field at line .*", "evaluation of \\[.*\\] failed, treating result as null", // TODO investigate? - // Awaiting fixes + // Awaiting fixes for query failure "Unknown column \\[\\]", // https://github.com/elastic/elasticsearch/issues/121741, "Plan \\[ProjectExec\\[\\[.* optimized incorrectly due to missing references", // https://github.com/elastic/elasticsearch/issues/125866 "optimized incorrectly due to missing references", // https://github.com/elastic/elasticsearch/issues/116781 - "The incoming YAML document exceeds the limit:" // still to investigate, but it seems to be specific to the test framework + "The incoming YAML document exceeds the limit:", // still to investigate, but it seems to be specific to the test framework + + // Awaiting fixes for correctness + "Expecting the following columns \\[.*\\], got" // https://github.com/elastic/elasticsearch/issues/129000 ); public static final Set ALLOWED_ERROR_PATTERNS = ALLOWED_ERRORS.stream() @@ -84,27 +88,73 @@ public void test() throws IOException { List indices = availableIndices(); List lookupIndices = lookupIndices(); List policies = availableEnrichPolicies(); + CommandGenerator.QuerySchema mappingInfo = new CommandGenerator.QuerySchema(indices, lookupIndices, policies); + EsqlQueryGenerator.QueryExecuted previousResult = null; for (int i = 0; i < ITERATIONS; i++) { - String command = EsqlQueryGenerator.sourceCommand(indices); + List previousCommands = new ArrayList<>(); + CommandGenerator commandGenerator = EsqlQueryGenerator.sourceCommand(); + CommandGenerator.CommandDescription desc = commandGenerator.generate(List.of(), List.of(), mappingInfo); + String command = desc.commandString(); EsqlQueryGenerator.QueryExecuted result = execute(command, 0); if (result.exception() != null) { checkException(result); - continue; + break; + } + if (checkResults(List.of(), commandGenerator, desc, null, result).success() == false) { + break; } + previousResult = result; + previousCommands.add(desc); for (int j = 0; j < MAX_DEPTH; j++) { if (result.outputSchema().isEmpty()) { break; } - command = EsqlQueryGenerator.pipeCommand(result.outputSchema(), policies, lookupIndices); + commandGenerator = EsqlQueryGenerator.randomPipeCommandGenerator(); + desc = commandGenerator.generate(previousCommands, result.outputSchema(), mappingInfo); + if (desc == CommandGenerator.EMPTY_DESCRIPTION) { + continue; + } + command = desc.commandString(); result = execute(result.query() + command, result.depth() + 1); if (result.exception() != null) { checkException(result); break; } + if (checkResults(previousCommands, commandGenerator, desc, previousResult, result).success() == false) { + break; + } + previousCommands.add(desc); + previousResult = result; } } } + private static CommandGenerator.ValidationResult checkResults( + List previousCommands, + CommandGenerator commandGenerator, + CommandGenerator.CommandDescription commandDescription, + EsqlQueryGenerator.QueryExecuted previousResult, + EsqlQueryGenerator.QueryExecuted result + ) { + CommandGenerator.ValidationResult outputValidation = commandGenerator.validateOutput( + previousCommands, + commandDescription, + previousResult == null ? null : previousResult.outputSchema(), + previousResult == null ? null : previousResult.result(), + result.outputSchema(), + result.result() + ); + if (outputValidation.success() == false) { + for (Pattern allowedError : ALLOWED_ERROR_PATTERNS) { + if (allowedError.matcher(outputValidation.errorMessage()).matches()) { + return outputValidation; + } + } + fail("query: " + result.query() + "\nerror: " + outputValidation.errorMessage()); + } + return outputValidation; + } + private void checkException(EsqlQueryGenerator.QueryExecuted query) { for (Pattern allowedError : ALLOWED_ERROR_PATTERNS) { if (allowedError.matcher(query.exception().getMessage()).matches()) { @@ -114,16 +164,18 @@ private void checkException(EsqlQueryGenerator.QueryExecuted query) { fail("query: " + query.query() + "\nexception: " + query.exception().getMessage()); } + @SuppressWarnings("unchecked") private EsqlQueryGenerator.QueryExecuted execute(String command, int depth) { try { Map a = RestEsqlTestCase.runEsqlSync(new RestEsqlTestCase.RequestObjectBuilder().query(command).build()); List outputSchema = outputSchema(a); - return new EsqlQueryGenerator.QueryExecuted(command, depth, outputSchema, null); + List> values = (List>) a.get("values"); + return new EsqlQueryGenerator.QueryExecuted(command, depth, outputSchema, values, null); } catch (Exception e) { - return new EsqlQueryGenerator.QueryExecuted(command, depth, null, e); + return new EsqlQueryGenerator.QueryExecuted(command, depth, null, null, e); } catch (AssertionError ae) { // this is for ensureNoWarnings() - return new EsqlQueryGenerator.QueryExecuted(command, depth, null, new RuntimeException(ae.getMessage())); + return new EsqlQueryGenerator.QueryExecuted(command, depth, null, null, new RuntimeException(ae.getMessage())); } } @@ -144,7 +196,7 @@ private List availableIndices() throws IOException { .toList(); } - record LookupIdx(String idxName, String key, String keyType) {} + public record LookupIdx(String idxName, String key, String keyType) {} private List lookupIndices() { List result = new ArrayList<>(); diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/README.asciidoc b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/README.asciidoc new file mode 100644 index 0000000000000..0c01e7c524ce3 --- /dev/null +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/README.asciidoc @@ -0,0 +1,43 @@ += ES|QL Generative Tests + +These tests generate random queries and execute them. + +The intention is not to test the single commands, but rather to test how ES|QL query engine +(parser, optimizers, query layout, compute) manages very complex queries. + +The test workflow is the following: + +1. Generate a source command (eg. `FROM idx`) +2. Execute it +3. Check the result +4. Based on the previous query output, generate a pipe command (eg. `| EVAL foo = to_lower(bar))` +5. Append the command to the query and execute it +6. Check the result +7. If the query is less than N commands (see `GenerativeRestTest.MAX_DEPTH)`, go to point `4` + +This workflow is executed M times (see `GenerativeRestTest.ITERATIONS`) + +The result check happens at two levels: + +* query success/failure - If the query fails: + ** If the error is in `GenerativeRestTest.ALLOWED_ERRORS`, ignore it and start with next iteration. + ** Otherwise throw an assertion error +* check result correctness - this is delegated to last executed command generator + +== 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. + +All you have to do is: + +* 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 diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/CommandGenerator.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/CommandGenerator.java new file mode 100644 index 0000000000000..7d652fddbc87a --- /dev/null +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/CommandGenerator.java @@ -0,0 +1,142 @@ +/* + * 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.rest.generative.command; + +import org.elasticsearch.xpack.esql.CsvTestsDataLoader; +import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator; +import org.elasticsearch.xpack.esql.qa.rest.generative.GenerativeRestTest; + +import java.util.List; +import java.util.Map; + +/** + * Implement this if you want to your command to be tested by the random query generator. + * Then add it to the right list in {@link EsqlQueryGenerator} + *

+ * The i + */ +public interface CommandGenerator { + + /** + * @param commandName the name of the command that is being generated + * @param commandString the full command string, including the "|" + * @param context additional information that could be useful for output validation. + * This will be passed to validateOutput after the query execution, together with the query output + */ + record CommandDescription(String commandName, CommandGenerator generator, String commandString, Map context) {} + + record QuerySchema( + List baseIndices, + List lookupIndices, + List enrichPolicies + ) {} + + record ValidationResult(boolean success, String errorMessage) {} + + CommandDescription EMPTY_DESCRIPTION = new CommandDescription("", new CommandGenerator() { + @Override + public CommandDescription generate( + List previousCommands, + List previousOutput, + QuerySchema schema + ) { + return EMPTY_DESCRIPTION; + } + + @Override + public ValidationResult validateOutput( + List previousCommands, + CommandDescription command, + List previousColumns, + List> previousOutput, + List columns, + List> output + ) { + return VALIDATION_OK; + } + }, "", Map.of()); + + ValidationResult VALIDATION_OK = new ValidationResult(true, null); + + /** + * Implement this method to generate a command, that will be appended to an existing query and then executed. + * See also {@link CommandDescription} + * + * @param previousCommands the list of the previous commands in the query + * @param previousOutput the output returned by the query so far. + * @param schema The columns returned by the query so far. It contains name and type information for each column. + * @return All the details about the generated command. See {@link CommandDescription}. + * If something goes wrong and for some reason you can't generate a command, you should return {@link CommandGenerator#EMPTY_DESCRIPTION} + */ + CommandDescription generate( + List previousCommands, + List previousOutput, + QuerySchema schema + ); + + /** + * This will be invoked after the query execution. + * You are expected to put validation logic in here. + * + * @param previousCommands The list of commands before the last generated one. + * @param command The description of the command you just generated. + * It also contains the context information you stored during command generation. + * @param previousColumns The output schema of the original query (without last generated command). + * It contains name and type information for each column, see {@link EsqlQueryGenerator.Column} + * @param previousOutput The output of the original query (without last generated command), as a list (rows) of lists (columns) of values + * @param columns The output schema of the full query (WITH last generated command). + * @param output The output of the full query (WITH last generated command), as a list (rows) of lists (columns) of values + * @return The result of the output validation. If the validation succeeds, you should return {@link CommandGenerator#VALIDATION_OK}. + * Also, if for some reason you can't validate the output, just return {@link CommandGenerator#VALIDATION_OK}; for a command, having a generator without + * validation is much better than having no generator at all. + */ + ValidationResult validateOutput( + List previousCommands, + CommandDescription command, + List previousColumns, + List> previousOutput, + List columns, + List> output + ); + + static ValidationResult expectSameRowCount( + List previousCommands, + List> previousOutput, + List> output + ) { + + // ES|QL is quite non-deterministic in this sense, we can't guarantee it for now + // if (output.size() != previousOutput.size()) { + // return new ValidationResult(false, "Expecting [" + previousOutput.size() + "] rows, but got [" + output.size() + "]"); + // } + + return VALIDATION_OK; + } + + static ValidationResult expectSameColumns(List previousColumns, List columns) { + + if (previousColumns.stream().anyMatch(x -> x.name().contains(""))) { + return VALIDATION_OK; // known bug + } + + if (previousColumns.size() != columns.size()) { + return new ValidationResult(false, "Expecting [" + previousColumns.size() + "] columns, got [" + columns.size() + "]"); + } + + List prevColNames = previousColumns.stream().map(EsqlQueryGenerator.Column::name).toList(); + List newColNames = columns.stream().map(EsqlQueryGenerator.Column::name).toList(); + if (prevColNames.equals(newColNames) == false) { + return new ValidationResult( + false, + "Expecting the following columns [" + String.join(", ", prevColNames) + "], got [" + String.join(", ", newColNames) + "]" + ); + } + + return VALIDATION_OK; + } +} diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/DissectGenerator.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/DissectGenerator.java new file mode 100644 index 0000000000000..83f8ae983dcde --- /dev/null +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/DissectGenerator.java @@ -0,0 +1,78 @@ +/* + * 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.rest.generative.command.pipe; + +import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator; +import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator; + +import java.util.List; +import java.util.Map; + +import static org.elasticsearch.test.ESTestCase.randomBoolean; +import static org.elasticsearch.test.ESTestCase.randomIntBetween; + +public class DissectGenerator implements CommandGenerator { + + public static final String DISSECT = "dissect"; + public static final CommandGenerator INSTANCE = new DissectGenerator(); + + @Override + public CommandDescription generate( + List previousCommands, + List previousOutput, + QuerySchema schema + ) { + String field = EsqlQueryGenerator.randomStringField(previousOutput); + if (field == null) { + return EMPTY_DESCRIPTION;// no strings to dissect, just skip + } + StringBuilder result = new StringBuilder(" | dissect "); + result.append(field); + result.append(" \""); + for (int i = 0; i < randomIntBetween(1, 3); i++) { + if (i > 0) { + result.append(" "); + } + result.append("%{"); + String fieldName; + if (randomBoolean()) { + fieldName = EsqlQueryGenerator.randomIdentifier(); + } else { + fieldName = EsqlQueryGenerator.randomRawName(previousOutput); + if (fieldName == null) { + fieldName = EsqlQueryGenerator.randomIdentifier(); + } + } + result.append(fieldName); + result.append("}"); + } + result.append("\""); + String cmdString = result.toString(); + return new CommandDescription(DISSECT, this, cmdString, Map.of()); + } + + @Override + public ValidationResult validateOutput( + List previousCommands, + CommandDescription commandDescription, + List previousColumns, + List> previousOutput, + List columns, + List> output + ) { + if (commandDescription == EMPTY_DESCRIPTION) { + return VALIDATION_OK; + } + + if (previousColumns.size() > columns.size()) { + return new ValidationResult(false, "Expecting at least [" + previousColumns.size() + "] columns, got [" + columns.size() + "]"); + } + + return CommandGenerator.expectSameRowCount(previousCommands, previousOutput, output); + } +} diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/DropGenerator.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/DropGenerator.java new file mode 100644 index 0000000000000..8bf2597f808c0 --- /dev/null +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/DropGenerator.java @@ -0,0 +1,91 @@ +/* + * 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.rest.generative.command.pipe; + +import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator; +import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator; + +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; + +import static org.elasticsearch.test.ESTestCase.randomBoolean; +import static org.elasticsearch.test.ESTestCase.randomIntBetween; + +public class DropGenerator implements CommandGenerator { + + public static final String DROP = "drop"; + public static final String DROPPED_COLUMNS = "dropped_columns"; + + public static final CommandGenerator INSTANCE = new DropGenerator(); + + @Override + public CommandDescription generate( + List previousCommands, + List previousOutput, + QuerySchema schema + ) { + if (previousOutput.size() < 2) { + return CommandGenerator.EMPTY_DESCRIPTION; // don't drop all of them, just do nothing + } + Set droppedColumns = new HashSet<>(); + int n = randomIntBetween(1, previousOutput.size() - 1); + Set proj = new HashSet<>(); + for (int i = 0; i < n; i++) { + String name = EsqlQueryGenerator.randomRawName(previousOutput); + if (name == null) { + continue; + } + if (name.length() > 1 && name.startsWith("`") == false && randomIntBetween(0, 100) < 10) { + if (randomBoolean()) { + name = name.substring(0, randomIntBetween(1, name.length() - 1)) + "*"; + } else { + name = "*" + name.substring(randomIntBetween(1, name.length() - 1)); + } + } else if (name.startsWith("`") == false && (randomBoolean() || name.isEmpty())) { + name = "`" + name + "`"; + } + proj.add(name); + droppedColumns.add(EsqlQueryGenerator.unquote(name)); + } + if (proj.isEmpty()) { + return CommandGenerator.EMPTY_DESCRIPTION; + } + String cmdString = " | drop " + proj.stream().collect(Collectors.joining(", ")); + return new CommandDescription(DROP, this, cmdString, Map.ofEntries(Map.entry(DROPPED_COLUMNS, droppedColumns))); + } + + @Override + @SuppressWarnings("unchecked") + public ValidationResult validateOutput( + List previousCommands, + CommandDescription commandDescription, + List previousColumns, + List> previousOutput, + List columns, + List> output + ) { + if (commandDescription == EMPTY_DESCRIPTION) { + return VALIDATION_OK; + } + Set droppedColumns = (Set) commandDescription.context().get(DROPPED_COLUMNS); + List resultColNames = columns.stream().map(EsqlQueryGenerator.Column::name).toList(); + // expected column names are unquoted already + for (String droppedColumn : droppedColumns) { + if (resultColNames.contains(droppedColumn)) { + return new ValidationResult(false, "Column [" + droppedColumn + "] was not dropped"); + } + } + // TODO awaits fix https://github.com/elastic/elasticsearch/issues/120272 + // return CommandGenerator.expectSameRowCount(previousOutput, output); + return VALIDATION_OK; + } + +} diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/EnrichGenerator.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/EnrichGenerator.java new file mode 100644 index 0000000000000..aac8f16e13285 --- /dev/null +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/EnrichGenerator.java @@ -0,0 +1,61 @@ +/* + * 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.rest.generative.command.pipe; + +import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator; +import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator; + +import java.util.List; +import java.util.Map; + +import static org.elasticsearch.test.ESTestCase.randomFrom; + +public class EnrichGenerator implements CommandGenerator { + + public static final String ENRICH = "enrich"; + public static final CommandGenerator INSTANCE = new EnrichGenerator(); + + @Override + public CommandDescription generate( + List previousCommands, + List previousOutput, + QuerySchema schema + ) { + String field = EsqlQueryGenerator.randomKeywordField(previousOutput); + if (field == null || schema.enrichPolicies().isEmpty()) { + return EMPTY_DESCRIPTION; + } + + // TODO add WITH + String cmdString = " | enrich " + + randomFrom(EsqlQueryGenerator.policiesOnKeyword(schema.enrichPolicies())).policyName() + + " on " + + field; + return new CommandDescription(ENRICH, this, cmdString, Map.of()); + } + + @Override + public ValidationResult validateOutput( + List previousCommands, + CommandDescription commandDescription, + List previousColumns, + List> previousOutput, + List columns, + List> output + ) { + if (commandDescription == EMPTY_DESCRIPTION) { + return VALIDATION_OK; + } + + if (previousColumns.size() > columns.size()) { + return new ValidationResult(false, "Expecting at least [" + previousColumns.size() + "] columns, got [" + columns.size() + "]"); + } + + return CommandGenerator.expectSameRowCount(previousCommands, previousOutput, output); + } +} diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/EvalGenerator.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/EvalGenerator.java new file mode 100644 index 0000000000000..b49e313fa4e3f --- /dev/null +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/EvalGenerator.java @@ -0,0 +1,96 @@ +/* + * 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.rest.generative.command.pipe; + +import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator; +import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +import static org.elasticsearch.test.ESTestCase.randomBoolean; +import static org.elasticsearch.test.ESTestCase.randomIntBetween; + +public class EvalGenerator implements CommandGenerator { + + public static final String EVAL = "eval"; + public static final String NEW_COLUMNS = "new_columns"; + public static final CommandGenerator INSTANCE = new EvalGenerator(); + + @Override + public CommandDescription generate( + List previousCommands, + List previousOutput, + QuerySchema schema + ) { + StringBuilder cmd = new StringBuilder(" | eval "); + int nFields = randomIntBetween(1, 10); + // TODO pass newly created fields to next expressions + var newColumns = new ArrayList<>(); + for (int i = 0; i < nFields; i++) { + String name; + if (randomBoolean()) { + name = EsqlQueryGenerator.randomIdentifier(); + } else { + name = EsqlQueryGenerator.randomName(previousOutput); + if (name == null) { + name = EsqlQueryGenerator.randomIdentifier(); + } + } + String expression = EsqlQueryGenerator.expression(previousOutput); + if (i > 0) { + cmd.append(","); + } + cmd.append(" "); + cmd.append(name); + newColumns.remove(unquote(name)); + newColumns.add(unquote(name)); + cmd.append(" = "); + cmd.append(expression); + } + String cmdString = cmd.toString(); + return new CommandDescription(EVAL, this, cmdString, Map.ofEntries(Map.entry(NEW_COLUMNS, newColumns))); + } + + @Override + @SuppressWarnings("unchecked") + public ValidationResult validateOutput( + List previousCommands, + CommandDescription commandDescription, + List previousColumns, + List> previousOutput, + List columns, + List> output + ) { + List expectedColumns = (List) commandDescription.context().get(NEW_COLUMNS); + List resultColNames = columns.stream().map(EsqlQueryGenerator.Column::name).toList(); + List lastColumns = resultColNames.subList(resultColNames.size() - expectedColumns.size(), resultColNames.size()); + lastColumns = lastColumns.stream().map(EvalGenerator::unquote).toList(); + // expected column names are unquoted already + if (columns.size() < expectedColumns.size() || lastColumns.equals(expectedColumns) == false) { + return new ValidationResult( + false, + "Expecting the following as last columns [" + + String.join(", ", expectedColumns) + + "] but got [" + + String.join(", ", resultColNames) + + "]" + ); + } + + return CommandGenerator.expectSameRowCount(previousCommands, previousOutput, output); + } + + private static String unquote(String colName) { + if (colName.startsWith("`") && colName.endsWith("`")) { + return colName.substring(1, colName.length() - 1); + } + return colName; + } +} diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/GrokGenerator.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/GrokGenerator.java new file mode 100644 index 0000000000000..60322eb12c351 --- /dev/null +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/GrokGenerator.java @@ -0,0 +1,75 @@ +/* + * 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.rest.generative.command.pipe; + +import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator; +import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator; + +import java.util.List; +import java.util.Map; + +import static org.elasticsearch.test.ESTestCase.randomBoolean; +import static org.elasticsearch.test.ESTestCase.randomIntBetween; + +public class GrokGenerator implements CommandGenerator { + + public static final String GROK = "grok"; + public static final CommandGenerator INSTANCE = new GrokGenerator(); + + @Override + public CommandDescription generate( + List previousCommands, + List previousOutput, + QuerySchema schema + ) { + String field = EsqlQueryGenerator.randomStringField(previousOutput); + if (field == null) { + return EMPTY_DESCRIPTION;// no strings to grok, just skip + } + StringBuilder result = new StringBuilder(" | grok "); + result.append(field); + result.append(" \""); + for (int i = 0; i < randomIntBetween(1, 3); i++) { + if (i > 0) { + result.append(" "); + } + result.append("%{WORD:"); + if (randomBoolean()) { + result.append(EsqlQueryGenerator.randomIdentifier()); + } else { + String fieldName = EsqlQueryGenerator.randomRawName(previousOutput); + if (fieldName == null) { + fieldName = EsqlQueryGenerator.randomIdentifier(); + } + result.append(fieldName); + } + result.append("}"); + } + result.append("\""); + String cmdString = result.toString(); + return new CommandDescription(GROK, this, cmdString, Map.of()); + } + + @Override + public ValidationResult validateOutput( + List previousCommands, + CommandDescription commandDescription, + List previousColumns, + List> previousOutput, + List columns, + List> output + ) { + if (commandDescription == EMPTY_DESCRIPTION) { + return VALIDATION_OK; + } + if (previousColumns.size() > columns.size()) { + return new ValidationResult(false, "Expecting at least [" + previousColumns.size() + "] columns, got [" + columns.size() + "]"); + } + return CommandGenerator.expectSameRowCount(previousCommands, previousOutput, output); + } +} diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/KeepGenerator.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/KeepGenerator.java new file mode 100644 index 0000000000000..f3f522576124e --- /dev/null +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/KeepGenerator.java @@ -0,0 +1,82 @@ +/* + * 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.rest.generative.command.pipe; + +import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator; +import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator; + +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; + +import static org.elasticsearch.test.ESTestCase.randomBoolean; +import static org.elasticsearch.test.ESTestCase.randomIntBetween; + +public class KeepGenerator implements CommandGenerator { + + public static final String KEEP = "keep"; + + public static final CommandGenerator INSTANCE = new KeepGenerator(); + + @Override + public CommandDescription generate( + List previousCommands, + List previousOutput, + QuerySchema schema + ) { + int n = randomIntBetween(1, previousOutput.size()); + Set proj = new HashSet<>(); + for (int i = 0; i < n; i++) { + if (randomIntBetween(0, 100) < 5) { + proj.add("*"); + } else { + String name = EsqlQueryGenerator.randomName(previousOutput); + if (name == null) { + continue; + } + if (name.length() > 1 && name.startsWith("`") == false && randomIntBetween(0, 100) < 10) { + if (randomBoolean()) { + name = name.substring(0, randomIntBetween(1, name.length() - 1)) + "*"; + } else { + name = "*" + name.substring(randomIntBetween(1, name.length() - 1)); + } + } + proj.add(name); + } + } + if (proj.isEmpty()) { + return EMPTY_DESCRIPTION; + } + String cmdString = " | keep " + proj.stream().collect(Collectors.joining(", ")); + return new CommandDescription(KEEP, this, cmdString, Map.of()); + } + + @Override + @SuppressWarnings("unchecked") + public ValidationResult validateOutput( + List previousCommands, + CommandDescription commandDescription, + List previousColumns, + List> previousOutput, + List columns, + List> output + ) { + if (commandDescription == EMPTY_DESCRIPTION) { + return VALIDATION_OK; + } + + if (previousColumns.size() < columns.size()) { + return new ValidationResult(false, "Expecting at most [" + previousColumns.size() + "] columns, got [" + columns.size() + "]"); + } + + return VALIDATION_OK; + } + +} diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/LimitGenerator.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/LimitGenerator.java new file mode 100644 index 0000000000000..d1e6fd8b3bfc8 --- /dev/null +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/LimitGenerator.java @@ -0,0 +1,57 @@ +/* + * 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.rest.generative.command.pipe; + +import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator; +import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator; + +import java.util.List; +import java.util.Map; + +import static org.elasticsearch.test.ESTestCase.randomIntBetween; + +public class LimitGenerator implements CommandGenerator { + + public static final String LIMIT = "limit"; + public static final CommandGenerator INSTANCE = new LimitGenerator(); + + @Override + public CommandDescription generate( + List previousCommands, + List previousOutput, + QuerySchema schema + ) { + int limit = randomIntBetween(0, 15000); + String cmd = " | limit " + limit; + return new CommandDescription(LIMIT, this, cmd, Map.ofEntries(Map.entry(LIMIT, limit))); + } + + @Override + @SuppressWarnings("unchecked") + public ValidationResult validateOutput( + List previousCommands, + CommandDescription commandDescription, + List previousColumns, + List> previousOutput, + List columns, + List> output + ) { + int limit = (int) commandDescription.context().get(LIMIT); + boolean defaultLimit = false; + for (CommandDescription previousCommand : previousCommands) { + if (previousCommand.commandName().equals(LIMIT)) { + defaultLimit = true; + } + } + + if (previousOutput.size() > limit && output.size() != limit || defaultLimit && previousOutput.size() < output.size()) { + return new ValidationResult(false, "Expecting [" + limit + "] records, got [" + output.size() + "]"); + } + return CommandGenerator.expectSameColumns(previousColumns, columns); + } +} diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/LookupJoinGenerator.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/LookupJoinGenerator.java new file mode 100644 index 0000000000000..4af6c5d73090d --- /dev/null +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/LookupJoinGenerator.java @@ -0,0 +1,63 @@ +/* + * 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.rest.generative.command.pipe; + +import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator; +import org.elasticsearch.xpack.esql.qa.rest.generative.GenerativeRestTest; +import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator; + +import java.util.List; +import java.util.Map; + +import static org.elasticsearch.test.ESTestCase.randomFrom; + +public class LookupJoinGenerator implements CommandGenerator { + + public static final String LOOKUP_JOIN = "lookup join"; + public static final CommandGenerator INSTANCE = new LookupJoinGenerator(); + + @Override + public CommandDescription generate( + List previousCommands, + List previousOutput, + QuerySchema schema + ) { + GenerativeRestTest.LookupIdx lookupIdx = randomFrom(schema.lookupIndices()); + String lookupIdxName = lookupIdx.idxName(); + String idxKey = lookupIdx.key(); + String keyType = lookupIdx.keyType(); + + var candidateKeys = previousOutput.stream().filter(x -> x.type().equals(keyType)).toList(); + if (candidateKeys.isEmpty()) { + return EMPTY_DESCRIPTION; + } + EsqlQueryGenerator.Column key = randomFrom(candidateKeys); + String cmdString = "| rename " + key.name() + " as " + idxKey + " | lookup join " + lookupIdxName + " on " + idxKey; + return new CommandDescription(LOOKUP_JOIN, this, cmdString, Map.of()); + } + + @Override + public ValidationResult validateOutput( + List previousCommands, + CommandDescription commandDescription, + List previousColumns, + List> previousOutput, + List columns, + List> output + ) { + if (commandDescription == EMPTY_DESCRIPTION) { + return VALIDATION_OK; + } + + // the -1 is for the additional RENAME, that could drop one column + if (previousColumns.size() - 1 > columns.size()) { + return new ValidationResult(false, "Expecting at least [" + previousColumns.size() + "] columns, got [" + columns.size() + "]"); + } + return VALIDATION_OK; + } +} diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/MvExpandGenerator.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/MvExpandGenerator.java new file mode 100644 index 0000000000000..317a2e459094e --- /dev/null +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/MvExpandGenerator.java @@ -0,0 +1,51 @@ +/* + * 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.rest.generative.command.pipe; + +import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator; +import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator; + +import java.util.List; +import java.util.Map; + +public class MvExpandGenerator implements CommandGenerator { + + public static final String MV_EXPAND = "mv_expand"; + + public static final CommandGenerator INSTANCE = new MvExpandGenerator(); + + @Override + public CommandDescription generate( + List previousCommands, + List previousOutput, + QuerySchema schema + ) { + String toExpand = EsqlQueryGenerator.randomName(previousOutput); + if (toExpand == null) { + return EMPTY_DESCRIPTION; // no columns to expand + } + String cmdString = " | mv_expand " + toExpand; + return new CommandDescription(MV_EXPAND, this, cmdString, Map.of()); + } + + @Override + public ValidationResult validateOutput( + List previousCommands, + CommandDescription commandDescription, + List previousColumns, + List> previousOutput, + List columns, + List> output + ) { + if (commandDescription == EMPTY_DESCRIPTION) { + return VALIDATION_OK; + } + return CommandGenerator.expectSameColumns(previousColumns, columns); + } + +} diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/RenameGenerator.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/RenameGenerator.java new file mode 100644 index 0000000000000..80b36c06e524e --- /dev/null +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/RenameGenerator.java @@ -0,0 +1,103 @@ +/* + * 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.rest.generative.command.pipe; + +import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator; +import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import static org.elasticsearch.test.ESTestCase.randomBoolean; +import static org.elasticsearch.test.ESTestCase.randomFrom; +import static org.elasticsearch.test.ESTestCase.randomIntBetween; + +public class RenameGenerator implements CommandGenerator { + + public static final String RENAME = "rename"; + + public static final CommandGenerator INSTANCE = new RenameGenerator(); + + @Override + public CommandDescription generate( + List previousCommands, + List previousOutput, + QuerySchema schema + ) { + int n = randomIntBetween(1, Math.min(3, previousOutput.size())); + List proj = new ArrayList<>(); + + Map nameToType = new HashMap<>(); + for (EsqlQueryGenerator.Column column : previousOutput) { + nameToType.put(column.name(), column.type()); + } + List names = new ArrayList<>( + previousOutput.stream() + .filter(EsqlQueryGenerator::fieldCanBeUsed) + .map(EsqlQueryGenerator.Column::name) + .collect(Collectors.toList()) + ); + if (names.isEmpty()) { + return EMPTY_DESCRIPTION; + } + for (int i = 0; i < n; i++) { + if (names.isEmpty()) { + break; + } + var name = randomFrom(names); + if (nameToType.get(name).endsWith("_range")) { + // ranges are not fully supported yet + continue; + } + names.remove(name); + + String newName; + if (names.isEmpty() || randomBoolean()) { + newName = EsqlQueryGenerator.randomIdentifier(); + names.add(newName); + } else { + newName = names.get(randomIntBetween(0, names.size() - 1)); + } + nameToType.put(newName, nameToType.get(name)); + if (randomBoolean() && name.startsWith("`") == false) { + name = "`" + name + "`"; + } + if (randomBoolean() && newName.startsWith("`") == false) { + newName = "`" + newName + "`"; + } + proj.add(name + " AS " + newName); + } + if (proj.isEmpty()) { + return EMPTY_DESCRIPTION; + } + String cmdString = " | rename " + proj.stream().collect(Collectors.joining(", ")); + return new CommandDescription(RENAME, this, cmdString, Map.of()); + } + + @Override + public ValidationResult validateOutput( + List previousCommands, + CommandDescription commandDescription, + List previousColumns, + List> previousOutput, + List columns, + List> output + ) { + if (commandDescription == EMPTY_DESCRIPTION) { + return VALIDATION_OK; + } + if (previousColumns.size() < columns.size()) { + return new ValidationResult(false, "Expecting at most [" + previousColumns.size() + "] columns, got [" + columns.size() + "]"); + } + return CommandGenerator.expectSameRowCount(previousCommands, previousOutput, output); + } + +} diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/SortGenerator.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/SortGenerator.java new file mode 100644 index 0000000000000..f7849d1c202f1 --- /dev/null +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/SortGenerator.java @@ -0,0 +1,60 @@ +/* + * 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.rest.generative.command.pipe; + +import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator; +import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator; + +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; + +import static org.elasticsearch.test.ESTestCase.randomFrom; +import static org.elasticsearch.test.ESTestCase.randomIntBetween; + +public class SortGenerator implements CommandGenerator { + + public static final String SORT = "sort"; + public static final CommandGenerator INSTANCE = new SortGenerator(); + + @Override + public CommandDescription generate( + List previousCommands, + List previousOutput, + QuerySchema schema + ) { + int n = randomIntBetween(1, previousOutput.size()); + Set proj = new HashSet<>(); + for (int i = 0; i < n; i++) { + String col = EsqlQueryGenerator.randomSortableName(previousOutput); + if (col == null) { + return EMPTY_DESCRIPTION; // no sortable columns + } + proj.add(col); + } + String cmd = " | sort " + + proj.stream() + .map(x -> x + randomFrom("", " ASC", " DESC") + randomFrom("", " NULLS FIRST", " NULLS LAST")) + .collect(Collectors.joining(", ")); + return new CommandDescription(SORT, this, cmd, Map.of()); + } + + @Override + public ValidationResult validateOutput( + List previousCommands, + CommandDescription commandDescription, + List previousColumns, + List> previousOutput, + List columns, + List> output + ) { + return CommandGenerator.expectSameColumns(previousColumns, columns); + } +} diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/StatsGenerator.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/StatsGenerator.java new file mode 100644 index 0000000000000..b0ce7f43af997 --- /dev/null +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/StatsGenerator.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.rest.generative.command.pipe; + +import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator; +import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator; + +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import static org.elasticsearch.test.ESTestCase.randomBoolean; +import static org.elasticsearch.test.ESTestCase.randomIntBetween; + +public class StatsGenerator implements CommandGenerator { + + public static final String STATS = "stats"; + public static final CommandGenerator INSTANCE = new StatsGenerator(); + + @Override + public CommandDescription generate( + List previousCommands, + List previousOutput, + QuerySchema schema + ) { + List nonNull = previousOutput.stream() + .filter(EsqlQueryGenerator::fieldCanBeUsed) + .filter(x -> x.type().equals("null") == false) + .collect(Collectors.toList()); + if (nonNull.isEmpty()) { + return EMPTY_DESCRIPTION; + } + StringBuilder cmd = new StringBuilder(" | stats "); + int nStats = randomIntBetween(1, 5); + for (int i = 0; i < nStats; i++) { + String name; + if (randomBoolean()) { + name = EsqlQueryGenerator.randomIdentifier(); + } else { + name = EsqlQueryGenerator.randomName(previousOutput); + if (name == null) { + name = EsqlQueryGenerator.randomIdentifier(); + } + } + String expression = EsqlQueryGenerator.agg(nonNull); + if (i > 0) { + cmd.append(","); + } + cmd.append(" "); + cmd.append(name); + cmd.append(" = "); + cmd.append(expression); + } + if (randomBoolean()) { + var col = EsqlQueryGenerator.randomGroupableName(nonNull); + if (col != null) { + cmd.append(" by " + col); + } + } + return new CommandDescription(STATS, this, cmd.toString(), Map.of()); + } + + @Override + public ValidationResult validateOutput( + List previousCommands, + CommandDescription commandDescription, + List previousColumns, + List> previousOutput, + List columns, + List> output + ) { + // TODO validate columns + return VALIDATION_OK; + } +} diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/WhereGenerator.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/WhereGenerator.java new file mode 100644 index 0000000000000..9bba468de0412 --- /dev/null +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/WhereGenerator.java @@ -0,0 +1,63 @@ +/* + * 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.rest.generative.command.pipe; + +import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator; +import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator; + +import java.util.List; +import java.util.Map; + +import static org.elasticsearch.test.ESTestCase.randomBoolean; +import static org.elasticsearch.test.ESTestCase.randomIntBetween; + +public class WhereGenerator implements CommandGenerator { + + public static final String WHERE = "where"; + public static final CommandGenerator INSTANCE = new WhereGenerator(); + + @Override + public CommandDescription generate( + List previousCommands, + List previousOutput, + QuerySchema schema + ) { + // TODO more complex conditions + StringBuilder result = new StringBuilder(" | where "); + int nConditions = randomIntBetween(1, 5); + for (int i = 0; i < nConditions; i++) { + String exp = EsqlQueryGenerator.booleanExpression(previousOutput); + if (exp == null) { + // cannot generate expressions, just skip + return EMPTY_DESCRIPTION; + } + if (i > 0) { + result.append(randomBoolean() ? " AND " : " OR "); + } + if (randomBoolean()) { + result.append(" NOT "); + } + result.append(exp); + } + + String cmd = result.toString(); + return new CommandDescription(WHERE, this, cmd, Map.of()); + } + + @Override + public ValidationResult validateOutput( + List previousCommands, + CommandDescription commandDescription, + List previousColumns, + List> previousOutput, + List columns, + List> output + ) { + return CommandGenerator.expectSameColumns(previousColumns, columns); + } +} diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/source/FromGenerator.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/source/FromGenerator.java new file mode 100644 index 0000000000000..05cd307a50755 --- /dev/null +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/source/FromGenerator.java @@ -0,0 +1,54 @@ +/* + * 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.rest.generative.command.source; + +import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator; +import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator; + +import java.util.List; +import java.util.Map; + +import static org.elasticsearch.test.ESTestCase.randomIntBetween; +import static org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator.indexPattern; + +public class FromGenerator implements CommandGenerator { + + public static final FromGenerator INSTANCE = new FromGenerator(); + + @Override + public CommandDescription generate( + List previousCommands, + List previousOutput, + QuerySchema schema + ) { + StringBuilder result = new StringBuilder("from "); + int items = randomIntBetween(1, 3); + List availableIndices = schema.baseIndices(); + for (int i = 0; i < items; i++) { + String pattern = indexPattern(availableIndices.get(randomIntBetween(0, availableIndices.size() - 1))); + if (i > 0) { + result.append(","); + } + result.append(pattern); + } + String query = result.toString(); + return new CommandDescription("from", this, query, Map.of()); + } + + @Override + public ValidationResult validateOutput( + List previousCommands, + CommandDescription commandDescription, + List previousColumns, + List> previousOutput, + List columns, + List> output + ) { + return VALIDATION_OK; + } +}