Skip to content

Commit 8a132d2

Browse files
felixbarnyfelixbarny
authored
Make PROMQL a source command (#138687)
1 parent 3324f0a commit 8a132d2

File tree

27 files changed

+2208
-1996
lines changed

27 files changed

+2208
-1996
lines changed

‎x-pack/plugin/esql/qa/server/multi-clusters/src/javaRestTest/java/org/elasticsearch/xpack/esql/ccq/MultiClusterSpecIT.java‎

Lines changed: 2 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import org.elasticsearch.xpack.esql.CsvSpecReader;
2525
import org.elasticsearch.xpack.esql.CsvSpecReader.CsvTestCase;
2626
import org.elasticsearch.xpack.esql.CsvTestsDataLoader;
27+
import org.elasticsearch.xpack.esql.EsqlTestUtils;
2728
import org.elasticsearch.xpack.esql.SpecReader;
2829
import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
2930
import org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase;
@@ -36,7 +37,6 @@
3637
import java.io.ByteArrayInputStream;
3738
import java.io.IOException;
3839
import java.net.URL;
39-
import java.util.Arrays;
4040
import java.util.List;
4141
import java.util.Locale;
4242
import java.util.Set;
@@ -299,43 +299,9 @@ static CsvSpecReader.CsvTestCase convertToRemoteIndices(CsvSpecReader.CsvTestCas
299299
dataLocation = randomFrom(DataLocation.values());
300300
}
301301
String query = testCase.query;
302-
String[] commands = query.split("\\|");
303-
String first = commands[0].trim();
304302
// If true, we're using *:index, otherwise we're using *:index,index
305303
boolean onlyRemotes = canUseRemoteIndicesOnly() && randomBoolean();
306-
307-
// Split "SET a=b; FROM x" into "SET a=b" and "FROM x"
308-
int lastSetDelimiterPosition = first.lastIndexOf(';');
309-
String setStatements = lastSetDelimiterPosition == -1 ? "" : first.substring(0, lastSetDelimiterPosition + 1);
310-
String afterSetStatements = lastSetDelimiterPosition == -1 ? first : first.substring(lastSetDelimiterPosition + 1);
311-
312-
// Split "FROM a, b, c" into "FROM" and "a, b, c"
313-
String[] commandParts = afterSetStatements.trim().split("\\s+", 2);
314-
315-
String command = commandParts[0].trim();
316-
if (command.equalsIgnoreCase("from") || command.equalsIgnoreCase("ts")) {
317-
String[] indexMetadataParts = commandParts[1].split("(?i)\\bmetadata\\b", 2);
318-
String indicesString = indexMetadataParts[0];
319-
String[] indices = indicesString.split(",");
320-
// This method may be called multiple times on the same testcase when using @Repeat
321-
boolean alreadyConverted = Arrays.stream(indices).anyMatch(i -> i.trim().startsWith("*:"));
322-
if (alreadyConverted == false) {
323-
if (Arrays.stream(indices).anyMatch(i -> LOOKUP_INDICES.contains(i.trim().toLowerCase(Locale.ROOT)))) {
324-
// If the query contains lookup indices, use only remotes to avoid duplication
325-
onlyRemotes = true;
326-
}
327-
final boolean onlyRemotesFinal = onlyRemotes;
328-
final String remoteIndices = Arrays.stream(indices)
329-
.map(index -> unquoteAndRequoteAsRemote(index.trim(), onlyRemotesFinal))
330-
.collect(Collectors.joining(","));
331-
String newFirstCommand = command
332-
+ " "
333-
+ remoteIndices
334-
+ " "
335-
+ (indexMetadataParts.length == 1 ? "" : "metadata " + indexMetadataParts[1]);
336-
testCase.query = setStatements + newFirstCommand + query.substring(first.length());
337-
}
338-
}
304+
testCase.query = EsqlTestUtils.addRemoteIndices(testCase.query, LOOKUP_INDICES, onlyRemotes);
339305

340306
int offset = testCase.query.length() - query.length();
341307
if (offset != 0) {
@@ -365,40 +331,6 @@ static boolean hasIndexMetadata(String query) {
365331
return false;
366332
}
367333

368-
/**
369-
* Since partial quoting is prohibited, we need to take the index name, unquote it,
370-
* convert it to a remote index, and then requote it. For example, "employees" is unquoted,
371-
* turned into the remote index *:employees, and then requoted to get "*:employees".
372-
* @param index Name of the index.
373-
* @param asRemoteIndexOnly If the return needs to be in the form of "*:idx,idx" or "*:idx".
374-
* @return A remote index pattern that's requoted.
375-
*/
376-
private static String unquoteAndRequoteAsRemote(String index, boolean asRemoteIndexOnly) {
377-
index = index.trim();
378-
379-
int numOfQuotes = 0;
380-
for (; numOfQuotes < index.length(); numOfQuotes++) {
381-
if (index.charAt(numOfQuotes) != '"') {
382-
break;
383-
}
384-
}
385-
386-
String unquoted = unquote(index, numOfQuotes);
387-
if (asRemoteIndexOnly) {
388-
return quote("*:" + unquoted, numOfQuotes);
389-
} else {
390-
return quote("*:" + unquoted + "," + unquoted, numOfQuotes);
391-
}
392-
}
393-
394-
private static String quote(String index, int numOfQuotes) {
395-
return "\"".repeat(numOfQuotes) + index + "\"".repeat(numOfQuotes);
396-
}
397-
398-
private static String unquote(String index, int numOfQuotes) {
399-
return index.substring(numOfQuotes, index.length() - numOfQuotes);
400-
}
401-
402334
@Override
403335
protected boolean enableRoundingDoubleValuesOnAsserting() {
404336
return true;

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ protected void shouldSkipTest(String testName) throws IOException {
6262

6363
assumeFalse(
6464
"Tests using PROMQL are not supported for now",
65-
testCase.requiredCapabilities.contains(PROMQL_PRE_TECH_PREVIEW_V2.capabilityName())
65+
testCase.requiredCapabilities.contains(PROMQL_PRE_TECH_PREVIEW_V3.capabilityName())
6666
);
6767

6868
assumeFalse(

‎x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,7 @@
109109
import org.elasticsearch.xpack.esql.inference.InferenceResolution;
110110
import org.elasticsearch.xpack.esql.inference.InferenceService;
111111
import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
112+
import org.elasticsearch.xpack.esql.parser.EsqlParser;
112113
import org.elasticsearch.xpack.esql.parser.QueryParam;
113114
import org.elasticsearch.xpack.esql.plan.EsqlStatement;
114115
import org.elasticsearch.xpack.esql.plan.IndexPattern;
@@ -118,6 +119,8 @@
118119
import org.elasticsearch.xpack.esql.plan.logical.Explain;
119120
import org.elasticsearch.xpack.esql.plan.logical.Limit;
120121
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
122+
import org.elasticsearch.xpack.esql.plan.logical.SourceCommand;
123+
import org.elasticsearch.xpack.esql.plan.logical.UnresolvedRelation;
121124
import org.elasticsearch.xpack.esql.plan.logical.local.EmptyLocalSupplier;
122125
import org.elasticsearch.xpack.esql.plan.logical.local.LocalRelation;
123126
import org.elasticsearch.xpack.esql.plan.logical.local.LocalSupplier;
@@ -168,6 +171,7 @@
168171
import java.util.TreeMap;
169172
import java.util.function.Predicate;
170173
import java.util.jar.JarInputStream;
174+
import java.util.stream.Collectors;
171175
import java.util.stream.DoubleStream;
172176
import java.util.stream.IntStream;
173177
import java.util.zip.ZipEntry;
@@ -1279,4 +1283,93 @@ private static PhysicalPlan ignoreIdsInPhysicalPlan(PhysicalPlan plan) {
12791283
return fragmentExec.withFragment(ignoredInFragment);
12801284
});
12811285
}
1286+
1287+
public static String addRemoteIndices(String query, Set<String> lookupIndices, boolean onlyRemotes) {
1288+
String[] commands = query.split("\\|");
1289+
// remove subqueries
1290+
String first = commands[0].split(",\\s+\\(")[0].trim();
1291+
1292+
// Split "SET a=b; FROM x" into "SET a=b" and "FROM x"
1293+
int lastSetDelimiterPosition = first.lastIndexOf(';');
1294+
String setStatements = lastSetDelimiterPosition == -1 ? "" : first.substring(0, lastSetDelimiterPosition + 1);
1295+
String afterSetStatements = lastSetDelimiterPosition == -1 ? first : first.substring(lastSetDelimiterPosition + 1);
1296+
1297+
// Split "FROM a, b, c" into "FROM" and "a, b, c"
1298+
String[] commandParts = afterSetStatements.trim().split("\\s+", 2);
1299+
1300+
String command = commandParts[0].trim();
1301+
if (SourceCommand.isSourceCommand(command)) {
1302+
String commandArgs = commandParts[1].trim();
1303+
String[] indices = new EsqlParser().createStatement(afterSetStatements)
1304+
.collect(UnresolvedRelation.class)
1305+
.getFirst()
1306+
.indexPattern()
1307+
.indexPattern()
1308+
.split(",");
1309+
// This method may be called multiple times on the same testcase when using @Repeat
1310+
boolean alreadyConverted = Arrays.stream(indices).anyMatch(i -> i.trim().startsWith("*:"));
1311+
if (alreadyConverted == false) {
1312+
String stuffAfterIndices = getStuffAfterIndices(commandArgs, indices[indices.length - 1]);
1313+
if (Arrays.stream(indices).anyMatch(i -> lookupIndices.contains(i.trim().toLowerCase(Locale.ROOT)))) {
1314+
// If the query contains lookup indices, use only remotes to avoid duplication
1315+
onlyRemotes = true;
1316+
}
1317+
final boolean onlyRemotesFinal = onlyRemotes;
1318+
final String remoteIndices = Arrays.stream(indices)
1319+
.map(index -> unquoteAndRequoteAsRemote(index.trim(), onlyRemotesFinal))
1320+
.collect(Collectors.joining(","));
1321+
String newFirstCommand = command + " " + remoteIndices + " " + stuffAfterIndices;
1322+
return (setStatements + " " + newFirstCommand.trim() + query.substring(first.length())).trim();
1323+
}
1324+
}
1325+
return query;
1326+
}
1327+
1328+
private static String getStuffAfterIndices(String commandArgs, String lastIndex) {
1329+
String stuffAfterIndices;
1330+
if (commandArgs.contains(lastIndex)) {
1331+
stuffAfterIndices = commandArgs.substring(commandArgs.lastIndexOf(lastIndex) + lastIndex.length()).trim();
1332+
} else {
1333+
stuffAfterIndices = commandArgs.trim();
1334+
}
1335+
while (stuffAfterIndices.startsWith("\"")) {
1336+
stuffAfterIndices = stuffAfterIndices.substring(1).trim();
1337+
}
1338+
return stuffAfterIndices;
1339+
}
1340+
1341+
/**
1342+
* Since partial quoting is prohibited, we need to take the index name, unquote it,
1343+
* convert it to a remote index, and then requote it. For example, "employees" is unquoted,
1344+
* turned into the remote index *:employees, and then requoted to get "*:employees".
1345+
* @param index Name of the index.
1346+
* @param asRemoteIndexOnly If the return needs to be in the form of "*:idx,idx" or "*:idx".
1347+
* @return A remote index pattern that's requoted.
1348+
*/
1349+
private static String unquoteAndRequoteAsRemote(String index, boolean asRemoteIndexOnly) {
1350+
index = index.trim();
1351+
1352+
int numOfQuotes = 0;
1353+
for (; numOfQuotes < index.length(); numOfQuotes++) {
1354+
if (index.charAt(numOfQuotes) != '"') {
1355+
break;
1356+
}
1357+
}
1358+
1359+
String unquoted = unquote(index, numOfQuotes);
1360+
if (asRemoteIndexOnly) {
1361+
return quote("*:" + unquoted, numOfQuotes);
1362+
} else {
1363+
return quote("*:" + unquoted + "," + unquoted, numOfQuotes);
1364+
}
1365+
}
1366+
1367+
private static String quote(String index, int numOfQuotes) {
1368+
return "\"".repeat(numOfQuotes) + index + "\"".repeat(numOfQuotes);
1369+
}
1370+
1371+
private static String unquote(String index, int numOfQuotes) {
1372+
return index.substring(numOfQuotes, index.length() - numOfQuotes);
1373+
}
1374+
12821375
}

‎x-pack/plugin/esql/qa/testFixtures/src/main/resources/k8s-timeseries-avg-over-time.csv-spec‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,8 @@ cost:double | time_bucket:datetime
1818
;
1919

2020
avg_over_time_of_double_no_grouping_promql
21-
required_capability: promql_pre_tech_preview_v2
22-
TS k8s
23-
| PROMQL step 1m (sum(avg_over_time(network.cost[1m])))
21+
required_capability: promql_pre_tech_preview_v3
22+
PROMQL k8s step 1m (sum(avg_over_time(network.cost[1m])))
2423
| SORT `sum(avg_over_time(network.cost[1m]))` DESC, step DESC | LIMIT 10;
2524

2625
sum(avg_over_time(network.cost[1m])):double | step:date
@@ -47,9 +46,8 @@ cost:double | time_bucket:datetime
4746
;
4847

4948
avg_over_time_of_double_no_grouping_single_bucket_promql
50-
required_capability: promql_pre_tech_preview_v2
51-
TS k8s
52-
| PROMQL step 1h (sum(avg_over_time(network.cost[1h])));
49+
required_capability: promql_pre_tech_preview_v3
50+
PROMQL k8s step 1h (sum(avg_over_time(network.cost[1h])));
5351

5452
sum(avg_over_time(network.cost[1h])):double | step:date
5553
56.32254035241995 | 2024-05-10T00:00:00.000Z

‎x-pack/plugin/esql/qa/testFixtures/src/main/resources/k8s-timeseries-promql.csv-spec‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,8 @@
11
promql_start_end_step
2-
required_capability: promql_pre_tech_preview_v2
3-
TS k8s
4-
| PROMQL step 5m start "2024-05-10T00:20:00.000Z" end "2024-05-10T00:25:00.000Z" (
5-
sum(avg_over_time(network.cost[5m]))
6-
);
2+
required_capability: promql_pre_tech_preview_v3
3+
PROMQL k8s step 5m start "2024-05-10T00:20:00.000Z" end "2024-05-10T00:25:00.000Z" (
4+
sum(avg_over_time(network.cost[5m]))
5+
);
76

87
sum(avg_over_time(network.cost[5m])):double | step:date
98
50.25 | 2024-05-10T00:20:00.000Z

‎x-pack/plugin/esql/qa/testFixtures/src/main/resources/k8s-timeseries.csv-spec‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -133,9 +133,8 @@ max(rate(network.total_bytes_in)): double | time_bucket:date
133133
;
134134

135135
oneRateWithPromql
136-
required_capability: promql_pre_tech_preview_v2
137-
TS k8s
138-
| PROMQL step 5m (max(rate(network.total_bytes_in[5m])))
136+
required_capability: promql_pre_tech_preview_v3
137+
PROMQL k8s step 5m (max(rate(network.total_bytes_in[5m])))
139138
| SORT step DESC | LIMIT 2;
140139

141140
max(rate(network.total_bytes_in[5m])):double | step:date
@@ -154,9 +153,8 @@ max(rate(network.total_bytes_in)): double | time_bucket:date
154153
;
155154

156155
oneRateWithSingleStepPromql
157-
required_capability: promql_pre_tech_preview_v2
158-
TS k8s
159-
| PROMQL step 1h (max(rate(network.total_bytes_in[1h])));
156+
required_capability: promql_pre_tech_preview_v3
157+
PROMQL k8s step 1h (max(rate(network.total_bytes_in[1h])));
160158

161159
max(rate(network.total_bytes_in[1h])):double | step:date
162160
4.379885057471264 | 2024-05-10T00:00:00.000Z

‎x-pack/plugin/esql/qa/testFixtures/src/test/java/org/elasticsearch/xpack/esql/EsqlTestUtilsTests.java‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
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;
9+
10+
import org.elasticsearch.test.ESTestCase;
11+
12+
import java.util.Set;
13+
14+
import static org.hamcrest.Matchers.equalTo;
15+
16+
public class EsqlTestUtilsTests extends ESTestCase {
17+
18+
public void testPromQL() {
19+
assertThat(
20+
EsqlTestUtils.addRemoteIndices("PROMQL foo, bar step 1m (avg(baz))", Set.of(), false),
21+
equalTo("PROMQL *:foo,foo,*:bar,bar step 1m (avg(baz))")
22+
);
23+
assertThat(
24+
EsqlTestUtils.addRemoteIndices("PROMQL \"foo\", \"bar\" step 1m (avg(baz))", Set.of(), false),
25+
equalTo("PROMQL *:foo,foo,*:bar,bar step 1m (avg(baz))")
26+
);
27+
}
28+
29+
public void testPromQLDefaultIndex() {
30+
assertThat(
31+
EsqlTestUtils.addRemoteIndices("PROMQL step 1m (avg(baz))", Set.of(), false),
32+
equalTo("PROMQL *:*,* step 1m (avg(baz))")
33+
);
34+
}
35+
36+
public void testSet() {
37+
assertThat(
38+
EsqlTestUtils.addRemoteIndices("SET a=b; FROM foo | SORT bar", Set.of(), false),
39+
equalTo("SET a=b; FROM *:foo,foo | SORT bar")
40+
);
41+
}
42+
43+
public void testMetadata() {
44+
assertThat(
45+
EsqlTestUtils.addRemoteIndices("FROM foo METADATA _source | SORT bar", Set.of(), false),
46+
equalTo("FROM *:foo,foo METADATA _source | SORT bar")
47+
);
48+
}
49+
50+
public void testTS() {
51+
assertThat(
52+
EsqlTestUtils.addRemoteIndices("TS foo, \"bar\",baz | SORT bar", Set.of(), false),
53+
equalTo("TS *:foo,foo,*:bar,bar,*:baz,baz | SORT bar")
54+
);
55+
}
56+
57+
public void testIndexPatternWildcard() {
58+
assertThat(EsqlTestUtils.addRemoteIndices("TS fo* | SORT bar", Set.of(), false), equalTo("TS *:fo*,fo* | SORT bar"));
59+
}
60+
61+
public void testDuplicateIndex() {
62+
assertThat(
63+
EsqlTestUtils.addRemoteIndices("TS foo,bar,foo | SORT bar", Set.of(), false),
64+
equalTo("TS *:foo,foo,*:bar,bar,*:foo,foo | SORT bar")
65+
);
66+
}
67+
68+
public void testSubquery() {
69+
assertThat(EsqlTestUtils.addRemoteIndices("""
70+
FROM employees, (FROM employees_incompatible
71+
| ENRICH languages_policy on languages with language_name )
72+
metadata _index
73+
| EVAL emp_no = emp_no::long
74+
| WHERE emp_no >= 10091 AND emp_no < 10094
75+
| SORT _index, emp_no
76+
| KEEP _index, emp_no, languages, language_name""", Set.of(), false), equalTo("""
77+
FROM *:employees,employees, (FROM employees_incompatible
78+
| ENRICH languages_policy on languages with language_name )
79+
metadata _index
80+
| EVAL emp_no = emp_no::long
81+
| WHERE emp_no >= 10091 AND emp_no < 10094
82+
| SORT _index, emp_no
83+
| KEEP _index, emp_no, languages, language_name"""));
84+
}
85+
86+
public void testTripleQuotes() {
87+
assertThat(
88+
EsqlTestUtils.addRemoteIndices("from \"\"\"employees\"\"\" | limit 2", Set.of(), false),
89+
equalTo("from *:employees,employees | limit 2")
90+
);
91+
}
92+
93+
public void testRow() {
94+
assertThat(EsqlTestUtils.addRemoteIndices("""
95+
ROW a = "1953-01-23T12:15:00Z - some text - 127.0.0.1;"\s
96+
| DISSECT a "%{Y}-%{M}-%{D}T%{h}:%{m}:%{s}Z - %{msg} - %{ip};"\s
97+
| KEEP Y, M, D, h, m, s, msg, ip""", Set.of(), false), equalTo("""
98+
ROW a = "1953-01-23T12:15:00Z - some text - 127.0.0.1;"\s
99+
| DISSECT a "%{Y}-%{M}-%{D}T%{h}:%{m}:%{s}Z - %{msg} - %{ip};"\s
100+
| KEEP Y, M, D, h, m, s, msg, ip"""));
101+
}
102+
}

0 commit comments

Comments
 (0)