Skip to content

Commit e64b6d4

Browse files
felixbarnyfelixbarny
committed
Properly replace remote indices in PROMQL command
1 parent 78e3f15 commit e64b6d4

File tree

4 files changed

+181
-72
lines changed

4 files changed

+181
-72
lines changed

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

Lines changed: 2 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -24,9 +24,9 @@
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;
29-
import org.elasticsearch.xpack.esql.plan.logical.SourceCommand;
3030
import org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase;
3131
import org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase;
3232
import org.junit.AfterClass;
@@ -37,7 +37,6 @@
3737
import java.io.ByteArrayInputStream;
3838
import java.io.IOException;
3939
import java.net.URL;
40-
import java.util.Arrays;
4140
import java.util.List;
4241
import java.util.Locale;
4342
import java.util.Set;
@@ -300,43 +299,9 @@ static CsvSpecReader.CsvTestCase convertToRemoteIndices(CsvSpecReader.CsvTestCas
300299
dataLocation = randomFrom(DataLocation.values());
301300
}
302301
String query = testCase.query;
303-
String[] commands = query.split("\\|");
304-
String first = commands[0].trim();
305302
// If true, we're using *:index, otherwise we're using *:index,index
306303
boolean onlyRemotes = canUseRemoteIndicesOnly() && randomBoolean();
307-
308-
// Split "SET a=b; FROM x" into "SET a=b" and "FROM x"
309-
int lastSetDelimiterPosition = first.lastIndexOf(';');
310-
String setStatements = lastSetDelimiterPosition == -1 ? "" : first.substring(0, lastSetDelimiterPosition + 1);
311-
String afterSetStatements = lastSetDelimiterPosition == -1 ? first : first.substring(lastSetDelimiterPosition + 1);
312-
313-
// Split "FROM a, b, c" into "FROM" and "a, b, c"
314-
String[] commandParts = afterSetStatements.trim().split("\\s+", 2);
315-
316-
String command = commandParts[0].trim();
317-
if (SourceCommand.isSourceCommand(command)) {
318-
String[] indexMetadataParts = commandParts[1].split("(?i)\\bmetadata\\b", 2);
319-
String indicesString = indexMetadataParts[0];
320-
String[] indices = indicesString.split(",");
321-
// This method may be called multiple times on the same testcase when using @Repeat
322-
boolean alreadyConverted = Arrays.stream(indices).anyMatch(i -> i.trim().startsWith("*:"));
323-
if (alreadyConverted == false) {
324-
if (Arrays.stream(indices).anyMatch(i -> LOOKUP_INDICES.contains(i.trim().toLowerCase(Locale.ROOT)))) {
325-
// If the query contains lookup indices, use only remotes to avoid duplication
326-
onlyRemotes = true;
327-
}
328-
final boolean onlyRemotesFinal = onlyRemotes;
329-
final String remoteIndices = Arrays.stream(indices)
330-
.map(index -> unquoteAndRequoteAsRemote(index.trim(), onlyRemotesFinal))
331-
.collect(Collectors.joining(","));
332-
String newFirstCommand = command
333-
+ " "
334-
+ remoteIndices
335-
+ " "
336-
+ (indexMetadataParts.length == 1 ? "" : "metadata " + indexMetadataParts[1]);
337-
testCase.query = setStatements + newFirstCommand + query.substring(first.length());
338-
}
339-
}
304+
testCase.query = EsqlTestUtils.addRemotes(testCase.query, LOOKUP_INDICES, onlyRemotes);
340305

341306
int offset = testCase.query.length() - query.length();
342307
if (offset != 0) {
@@ -366,40 +331,6 @@ static boolean hasIndexMetadata(String query) {
366331
return false;
367332
}
368333

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

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

Lines changed: 92 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,92 @@ private static PhysicalPlan ignoreIdsInPhysicalPlan(PhysicalPlan plan) {
12791283
return fragmentExec.withFragment(ignoredInFragment);
12801284
});
12811285
}
1286+
1287+
public static String addRemotes(String query, Set<String> lookupIndices, boolean onlyRemotes) {
1288+
String[] commands = query.split("\\|");
1289+
String first = commands[0].trim();
1290+
1291+
// Split "SET a=b; FROM x" into "SET a=b" and "FROM x"
1292+
int lastSetDelimiterPosition = first.lastIndexOf(';');
1293+
String setStatements = lastSetDelimiterPosition == -1 ? "" : first.substring(0, lastSetDelimiterPosition + 1);
1294+
String afterSetStatements = lastSetDelimiterPosition == -1 ? first : first.substring(lastSetDelimiterPosition + 1);
1295+
1296+
// Split "FROM a, b, c" into "FROM" and "a, b, c"
1297+
String[] commandParts = afterSetStatements.trim().split("\\s+", 2);
1298+
1299+
String command = commandParts[0].trim();
1300+
String commandArgs = commandParts[1].trim();
1301+
if (SourceCommand.isSourceCommand(command)) {
1302+
String[] indices = new EsqlParser().createStatement(afterSetStatements)
1303+
.collect(UnresolvedRelation.class)
1304+
.getFirst()
1305+
.indexPattern()
1306+
.indexPattern()
1307+
.split(",");
1308+
// This method may be called multiple times on the same testcase when using @Repeat
1309+
boolean alreadyConverted = Arrays.stream(indices).anyMatch(i -> i.trim().startsWith("*:"));
1310+
if (alreadyConverted == false) {
1311+
String stuffAfterIndices = getStuffAfterIndices(commandArgs, indices[indices.length - 1]);
1312+
if (Arrays.stream(indices).anyMatch(i -> lookupIndices.contains(i.trim().toLowerCase(Locale.ROOT)))) {
1313+
// If the query contains lookup indices, use only remotes to avoid duplication
1314+
onlyRemotes = true;
1315+
}
1316+
final boolean onlyRemotesFinal = onlyRemotes;
1317+
final String remoteIndices = Arrays.stream(indices)
1318+
.map(index -> unquoteAndRequoteAsRemote(index.trim(), onlyRemotesFinal))
1319+
.collect(Collectors.joining(","));
1320+
String newFirstCommand = command + " " + remoteIndices + " " + stuffAfterIndices;
1321+
return (setStatements + " " + newFirstCommand.trim() + query.substring(first.length())).trim();
1322+
}
1323+
}
1324+
return query;
1325+
}
1326+
1327+
private static String getStuffAfterIndices(String commandArgs, String lastIndex) {
1328+
String stuffAfterIndices;
1329+
if (commandArgs.contains(lastIndex)) {
1330+
stuffAfterIndices = commandArgs.substring(commandArgs.lastIndexOf(lastIndex) + lastIndex.length()).trim();
1331+
} else {
1332+
stuffAfterIndices = commandArgs.trim();
1333+
}
1334+
if (stuffAfterIndices.startsWith("\"")) {
1335+
stuffAfterIndices = stuffAfterIndices.substring(1).trim();
1336+
}
1337+
return stuffAfterIndices;
1338+
}
1339+
1340+
/**
1341+
* Since partial quoting is prohibited, we need to take the index name, unquote it,
1342+
* convert it to a remote index, and then requote it. For example, "employees" is unquoted,
1343+
* turned into the remote index *:employees, and then requoted to get "*:employees".
1344+
* @param index Name of the index.
1345+
* @param asRemoteIndexOnly If the return needs to be in the form of "*:idx,idx" or "*:idx".
1346+
* @return A remote index pattern that's requoted.
1347+
*/
1348+
private static String unquoteAndRequoteAsRemote(String index, boolean asRemoteIndexOnly) {
1349+
index = index.trim();
1350+
1351+
int numOfQuotes = 0;
1352+
for (; numOfQuotes < index.length(); numOfQuotes++) {
1353+
if (index.charAt(numOfQuotes) != '"') {
1354+
break;
1355+
}
1356+
}
1357+
1358+
String unquoted = unquote(index, numOfQuotes);
1359+
if (asRemoteIndexOnly) {
1360+
return quote("*:" + unquoted, numOfQuotes);
1361+
} else {
1362+
return quote("*:" + unquoted + "," + unquoted, numOfQuotes);
1363+
}
1364+
}
1365+
1366+
private static String quote(String index, int numOfQuotes) {
1367+
return "\"".repeat(numOfQuotes) + index + "\"".repeat(numOfQuotes);
1368+
}
1369+
1370+
private static String unquote(String index, int numOfQuotes) {
1371+
return index.substring(numOfQuotes, index.length() - numOfQuotes);
1372+
}
1373+
12821374
}

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

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
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.addRemotes("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.addRemotes("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.addRemotes("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.addRemotes("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.addRemotes("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.addRemotes("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(
59+
EsqlTestUtils.addRemotes("TS fo* | SORT bar", Set.of(), false),
60+
equalTo("TS *:fo*,fo* | SORT bar")
61+
);
62+
}
63+
64+
public void testDuplicateIndex() {
65+
assertThat(
66+
EsqlTestUtils.addRemotes("TS foo,bar,foo | SORT bar", Set.of(), false),
67+
equalTo("TS *:foo,foo,*:bar,bar,*:foo,foo | SORT bar")
68+
);
69+
}
70+
}

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/parser/promql/PromqlParamsTests.java renamed to x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/parser/promql/PromqlParserTests.java

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import org.elasticsearch.xpack.esql.parser.EsqlParser;
1414
import org.elasticsearch.xpack.esql.parser.ParsingException;
1515
import org.elasticsearch.xpack.esql.parser.QueryParams;
16+
import org.elasticsearch.xpack.esql.plan.logical.UnresolvedRelation;
1617
import org.elasticsearch.xpack.esql.plan.logical.promql.PromqlCommand;
1718
import org.junit.BeforeClass;
1819

@@ -25,10 +26,11 @@
2526
import static org.elasticsearch.xpack.esql.EsqlTestUtils.withDefaultLimitWarning;
2627
import static org.hamcrest.Matchers.containsString;
2728
import static org.hamcrest.Matchers.equalTo;
29+
import static org.hamcrest.Matchers.hasSize;
2830
import static org.hamcrest.Matchers.nullValue;
2931
import static org.junit.Assume.assumeTrue;
3032

31-
public class PromqlParamsTests extends ESTestCase {
33+
public class PromqlParserTests extends ESTestCase {
3234

3335
private static final EsqlParser parser = new EsqlParser();
3436

@@ -186,6 +188,20 @@ public void testRangeQueryMissingStep() {
186188
assertThat(e.getMessage(), containsString("Parameter [step] or [time] is required"));
187189
}
188190

191+
public void testParseMultipleIndices() {
192+
PromqlCommand promqlCommand = parse("PROMQL foo, bar step 5m (avg(foo))");
193+
List<UnresolvedRelation> unresolvedRelations = promqlCommand.collect(UnresolvedRelation.class);
194+
assertThat(unresolvedRelations, hasSize(1));
195+
assertThat(unresolvedRelations.getFirst().indexPattern().indexPattern(), equalTo("foo,bar"));
196+
}
197+
198+
public void testParseRemoteIndices() {
199+
PromqlCommand promqlCommand = parse("PROMQL *:foo,foo step 5m (avg(foo))");
200+
List<UnresolvedRelation> unresolvedRelations = promqlCommand.collect(UnresolvedRelation.class);
201+
assertThat(unresolvedRelations, hasSize(1));
202+
assertThat(unresolvedRelations.getFirst().indexPattern().indexPattern(), equalTo("*:foo,foo"));
203+
}
204+
189205
private static PromqlCommand parse(String query) {
190206
return as(parser.createStatement(query), PromqlCommand.class);
191207
}

0 commit comments

Comments
 (0)