Skip to content

Commit 4680aa2

Browse files
pabloempabloem
committed
Prototype (not working) of a fully generative ESQL test for Timeseries
1 parent 83f9510 commit 4680aa2

File tree

5 files changed

+454
-0
lines changed

5 files changed

+454
-0
lines changed

‎x-pack/plugin/esql/qa/server/build.gradle‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,4 +10,8 @@ dependencies {
1010
// Requirement for some ESQL-specific utilities
1111
implementation project(':x-pack:plugin:esql')
1212
api project(xpackModule('esql:qa:testFixtures'))
13+
implementation("org.testcontainers:testcontainers:${versions.testcontainer}")
14+
runtimeOnly "com.github.docker-java:docker-java-transport-zerodep:${versions.dockerJava}"
15+
runtimeOnly "com.github.docker-java:docker-java-transport:${versions.dockerJava}"
16+
runtimeOnly "com.github.docker-java:docker-java-core:${versions.dockerJava}"
1317
}

‎x-pack/plugin/esql/qa/server/multi-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/multi_node/BasicGenerativeTimeseriesRestIT.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.qa.multi_node;
9+
10+
import com.carrotsearch.randomizedtesting.annotations.ThreadLeakFilters;
11+
12+
import org.elasticsearch.test.TestClustersThreadFilter;
13+
import org.elasticsearch.test.cluster.ElasticsearchCluster;
14+
import org.elasticsearch.xpack.esql.qa.rest.generative.GenerativeTimeseriesRestTest;
15+
import org.junit.ClassRule;
16+
17+
@ThreadLeakFilters(filters = TestClustersThreadFilter.class)
18+
public class BasicGenerativeTimeseriesRestIT extends GenerativeTimeseriesRestTest {
19+
@ClassRule
20+
public static ElasticsearchCluster cluster = Clusters.testCluster(spec -> spec.plugin("inference-service-test"));
21+
22+
@Override
23+
protected String getTestRestCluster() {
24+
return cluster.getHttpAddresses();
25+
}
26+
}

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

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.qa.single_node;
9+
10+
import org.elasticsearch.action.bulk.BulkRequest;
11+
import org.elasticsearch.action.index.IndexRequest;
12+
import org.elasticsearch.common.settings.Settings;
13+
import org.elasticsearch.datageneration.DocumentGenerator;
14+
import org.elasticsearch.datageneration.DataGeneratorSpecification;
15+
import org.elasticsearch.datageneration.FieldType;
16+
import org.elasticsearch.datageneration.MappingGenerator;
17+
import org.elasticsearch.datageneration.TemplateGenerator;
18+
import org.elasticsearch.datageneration.fields.PredefinedField;
19+
import org.elasticsearch.test.ESTestCase;
20+
import org.elasticsearch.xcontent.XContentType;
21+
import org.elasticsearch.xpack.esql.action.EsqlQueryResponse;
22+
import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator;
23+
import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator;
24+
import org.junit.Test;
25+
26+
import java.util.ArrayList;
27+
import java.util.List;
28+
import java.util.Map;
29+
import java.util.stream.Collectors;
30+
31+
public class TimeseriesRandomizedTestCaseIT extends GenerativeIT {
32+
33+
private static final int MAX_DEPTH = 15;
34+
35+
@Test
36+
public void testRandomTimeseriesQueries() throws Exception {
37+
String indexName = "timeseries_test";
38+
generateAndIndexData(indexName, 200);
39+
40+
// Ten random queries
41+
for (int i = 0; i < 10; i++) {
42+
runRandomQuery(indexName);
43+
}
44+
}
45+
46+
private void runRandomQuery(String indexName) {
47+
CommandGenerator.QuerySchema mappingInfo = new CommandGenerator.QuerySchema(List.of(indexName), List.of(), List.of());
48+
List<CommandGenerator.CommandDescription> previousCommands = new ArrayList<>();
49+
50+
// Base query to select recent data
51+
String baseQuery = "FROM " + indexName + " | WHERE @timestamp > NOW() - INTERVAL 1 DAY";
52+
EsqlQueryResponse response = run(baseQuery);
53+
List<EsqlQueryGenerator.Column> currentSchema = toGeneratorSchema(response.columns());
54+
55+
String currentQuery = baseQuery;
56+
57+
for (int j = 0; j < MAX_DEPTH; j++) {
58+
if (currentSchema.isEmpty()) {
59+
break;
60+
}
61+
62+
CommandGenerator commandGenerator = EsqlQueryGenerator.randomPipeCommandGenerator();
63+
CommandGenerator.CommandDescription desc = commandGenerator.generate(previousCommands, currentSchema, mappingInfo);
64+
65+
if (desc == CommandGenerator.EMPTY_DESCRIPTION) {
66+
continue;
67+
}
68+
69+
String nextCommand = desc.commandString();
70+
currentQuery += nextCommand;
71+
72+
try {
73+
response = run(currentQuery);
74+
currentSchema = toGeneratorSchema(response.columns());
75+
previousCommands.add(desc);
76+
} catch (Exception e) {
77+
// For now, we fail on any exception. More sophisticated error handling could be added here.
78+
throw new AssertionError("Query failed: " + currentQuery, e);
79+
}
80+
}
81+
}
82+
83+
private DataGeneratorSpecification createDataSpec() {
84+
return DataGeneratorSpecification.builder()
85+
.withMaxFieldCountPerLevel(10)
86+
.withMaxObjectDepth(2)
87+
.withPredefinedFields(List.of(new PredefinedField.WithGenerator(
88+
"timestamp",
89+
FieldType.DATE,
90+
Map.of("type", "date"),
91+
(ignored) -> ESTestCase.randomPositiveTimeValue()
92+
)))
93+
.build();
94+
}
95+
96+
private void generateAndIndexData(String indexName, int numDocs) throws Exception {
97+
var spec = createDataSpec();
98+
DocumentGenerator dataGenerator = new DocumentGenerator(spec);
99+
var templateGen = new TemplateGenerator(spec).generate();
100+
var mappingGen = new MappingGenerator(spec).generate(templateGen);
101+
// Create a new index with a randomized mapping
102+
createIndex(indexName, Settings.EMPTY);
103+
104+
BulkRequest bulkRequest = new BulkRequest(indexName);
105+
for (int i = 0; i < numDocs; i++) {
106+
Map<String, Object> document = dataGenerator.generate(templateGen, mappingGen);
107+
bulkRequest.add(new IndexRequest().source(document, XContentType.JSON));
108+
}
109+
client().bulk(bulkRequest).actionGet();
110+
refresh(indexName);
111+
}
112+
113+
private List<EsqlQueryGenerator.Column> toGeneratorSchema(List<ColumnInfoImpl> columnInfos) {
114+
if (columnInfos == null) {
115+
return List.of();
116+
}
117+
return columnInfos.stream()
118+
.map(ci -> new EsqlQueryGenerator.Column(ci.name(), ci.outputType()))
119+
.collect(Collectors.toList());
120+
}
121+
}

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

Lines changed: 206 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,206 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.qa.rest.generative;
9+
10+
import org.elasticsearch.action.admin.indices.create.CreateIndexRequestBuilder;
11+
import org.elasticsearch.client.Request;
12+
import org.elasticsearch.client.ResponseException;
13+
import org.elasticsearch.common.Strings;
14+
import org.elasticsearch.test.rest.ESRestTestCase;
15+
import org.elasticsearch.xcontent.XContentBuilder;
16+
import org.elasticsearch.xcontent.XContentFactory;
17+
import org.elasticsearch.xpack.esql.AssertWarnings;
18+
import org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase;
19+
import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator;
20+
import org.junit.AfterClass;
21+
import org.junit.Before;
22+
import org.testcontainers.containers.GenericContainer;
23+
import org.testcontainers.utility.DockerImageName;
24+
25+
import java.io.IOException;
26+
import java.util.ArrayList;
27+
import java.util.List;
28+
import java.util.Map;
29+
import java.util.Set;
30+
import java.util.regex.Pattern;
31+
import java.util.stream.Collectors;
32+
33+
public abstract class GenerativeTimeseriesRestTest extends ESRestTestCase {
34+
35+
public static final int ITERATIONS = 200; // More test cases
36+
public static final int MAX_DEPTH = 8; // Fewer statements
37+
38+
private static final String PROMETHEUS_IMAGE = "prom/prometheus:latest";
39+
private static GenericContainer<?> prometheusContainer;
40+
41+
public static final Set<String> ALLOWED_ERRORS = Set.of(
42+
// "Reference \\[.*\\] is ambiguous",
43+
// "Cannot use field \\[.*\\] due to ambiguities",
44+
// "cannot sort on .*",
45+
// "argument of \\[count.*\\] must",
46+
// "Cannot use field \\[.*\\] with unsupported type \\[.*_range\\]",
47+
// "Unbounded sort not supported yet",
48+
// "The field names are too complex to process",
49+
// "must be \\[any type except counter types\\]",
50+
// "Unknown column \\[<all-fields-projected>\\]",
51+
// "Plan \\[ProjectExec\\[\\[<no-fields>.* optimized incorrectly due to missing references",
52+
// "optimized incorrectly due to missing references",
53+
// "The incoming YAML document exceeds the limit:",
54+
// "Data too large",
55+
// "Expecting the following columns \\[.*\\], got",
56+
// "Expecting at most \\[.*\\] columns, got \\[.*\\]"
57+
);
58+
59+
public static final Set<Pattern> ALLOWED_ERROR_PATTERNS = ALLOWED_ERRORS.stream()
60+
.map(x -> ".*" + x + ".*")
61+
.map(x -> Pattern.compile(x, Pattern.DOTALL))
62+
.collect(Collectors.toSet());
63+
64+
@Before
65+
public void setup() throws IOException {
66+
if (prometheusContainer == null) {
67+
// prometheusContainer = new GenericContainer<>(DockerImageName.parse(PROMETHEUS_IMAGE)).withExposedPorts(9090);
68+
// prometheusContainer.start();
69+
}
70+
71+
if (indexExists("timeseries") == false) {
72+
TimeseriesDataGenerationHelper dataGen = new TimeseriesDataGenerationHelper();
73+
Request createIndex = new Request("PUT", "/timeseries");
74+
XContentBuilder builder = XContentFactory.jsonBuilder().map(dataGen.getMapping().raw());
75+
createIndex.setJsonEntity("{\"mappings\": " + Strings.toString(builder) + "}");
76+
client().performRequest(createIndex);
77+
78+
for (int i = 0; i < 1000; i++) {
79+
Request indexRequest = new Request("POST", "/timeseries/_doc");
80+
indexRequest.setJsonEntity(dataGen.generateDocumentAsString());
81+
client().performRequest(indexRequest);
82+
}
83+
client().performRequest(new Request("POST", "/timeseries/_refresh"));
84+
}
85+
}
86+
87+
@AfterClass
88+
public static void wipeTestData() throws IOException {
89+
if (prometheusContainer != null) {
90+
prometheusContainer.stop();
91+
prometheusContainer = null;
92+
}
93+
try {
94+
adminClient().performRequest(new Request("DELETE", "/*"));
95+
} catch (ResponseException e) {
96+
if (e.getResponse().getStatusLine().getStatusCode() != 404) {
97+
throw e;
98+
}
99+
}
100+
}
101+
102+
public void test() throws IOException {
103+
List<String> indices = List.of("timeseries");
104+
CommandGenerator.QuerySchema mappingInfo = new CommandGenerator.QuerySchema(indices, List.of(), List.of());
105+
EsqlQueryGenerator.QueryExecuted previousResult = null;
106+
for (int i = 0; i < ITERATIONS; i++) {
107+
List<CommandGenerator.CommandDescription> previousCommands = new ArrayList<>();
108+
CommandGenerator commandGenerator = EsqlQueryGenerator.sourceCommand();
109+
CommandGenerator.CommandDescription desc = commandGenerator.generate(List.of(), List.of(), mappingInfo);
110+
String command = desc.commandString();
111+
EsqlQueryGenerator.QueryExecuted result = execute(command, 0);
112+
if (result.exception() != null) {
113+
checkException(result);
114+
continue;
115+
}
116+
if (checkResults(List.of(), commandGenerator, desc, null, result).success() == false) {
117+
continue;
118+
}
119+
previousResult = result;
120+
previousCommands.add(desc);
121+
for (int j = 0; j < MAX_DEPTH; j++) {
122+
if (result.outputSchema().isEmpty()) {
123+
break;
124+
}
125+
commandGenerator = EsqlQueryGenerator.randomPipeCommandGenerator();
126+
desc = commandGenerator.generate(previousCommands, result.outputSchema(), mappingInfo);
127+
if (desc == CommandGenerator.EMPTY_DESCRIPTION) {
128+
continue;
129+
}
130+
command = desc.commandString();
131+
result = execute(result.query() + command, result.depth() + 1);
132+
if (result.exception() != null) {
133+
checkException(result);
134+
break;
135+
}
136+
if (checkResults(previousCommands, commandGenerator, desc, previousResult, result).success() == false) {
137+
break;
138+
}
139+
previousCommands.add(desc);
140+
previousResult = result;
141+
}
142+
}
143+
}
144+
145+
private static CommandGenerator.ValidationResult checkResults(
146+
List<CommandGenerator.CommandDescription> previousCommands,
147+
CommandGenerator commandGenerator,
148+
CommandGenerator.CommandDescription commandDescription,
149+
EsqlQueryGenerator.QueryExecuted previousResult,
150+
EsqlQueryGenerator.QueryExecuted result
151+
) {
152+
CommandGenerator.ValidationResult outputValidation = commandGenerator.validateOutput(
153+
previousCommands,
154+
commandDescription,
155+
previousResult == null ? null : previousResult.outputSchema(),
156+
previousResult == null ? null : previousResult.result(),
157+
result.outputSchema(),
158+
result.result()
159+
);
160+
if (outputValidation.success() == false) {
161+
for (Pattern allowedError : ALLOWED_ERROR_PATTERNS) {
162+
if (allowedError.matcher(outputValidation.errorMessage()).matches()) {
163+
return outputValidation;
164+
}
165+
}
166+
fail("query: " + result.query() + "\nerror: " + outputValidation.errorMessage());
167+
}
168+
return outputValidation;
169+
}
170+
171+
private void checkException(EsqlQueryGenerator.QueryExecuted query) {
172+
for (Pattern allowedError : ALLOWED_ERROR_PATTERNS) {
173+
if (allowedError.matcher(query.exception().getMessage()).matches()) {
174+
return;
175+
}
176+
}
177+
fail("query: " + query.query() + "\nexception: " + query.exception().getMessage());
178+
}
179+
180+
@SuppressWarnings("unchecked")
181+
private EsqlQueryGenerator.QueryExecuted execute(String command, int depth) {
182+
try {
183+
Map<String, Object> a = RestEsqlTestCase.runEsql(
184+
new RestEsqlTestCase.RequestObjectBuilder().query(command).build(),
185+
new AssertWarnings.AllowedRegexes(List.of(Pattern.compile(".*"))),
186+
RestEsqlTestCase.Mode.SYNC
187+
);
188+
List<EsqlQueryGenerator.Column> outputSchema = outputSchema(a);
189+
List<List<Object>> values = (List<List<Object>>) a.get("values");
190+
return new EsqlQueryGenerator.QueryExecuted(command, depth, outputSchema, values, null);
191+
} catch (Exception e) {
192+
return new EsqlQueryGenerator.QueryExecuted(command, depth, null, null, e);
193+
} catch (AssertionError ae) {
194+
return new EsqlQueryGenerator.QueryExecuted(command, depth, null, null, new RuntimeException(ae.getMessage()));
195+
}
196+
}
197+
198+
@SuppressWarnings("unchecked")
199+
private List<EsqlQueryGenerator.Column> outputSchema(Map<String, Object> a) {
200+
List<Map<String, String>> cols = (List<Map<String, String>>) a.get("columns");
201+
if (cols == null) {
202+
return null;
203+
}
204+
return cols.stream().map(x -> new EsqlQueryGenerator.Column(x.get("name"), x.get("type"))).collect(Collectors.toList());
205+
}
206+
}

0 commit comments

Comments
 (0)