Skip to content

Commit d3bb902

Browse files
vamsimanoharvamsimanohar
authored
back quote fix (opensearch-project#1041)
Signed-off-by: vamsi-amazon <reddyvam@amazon.com>
1 parent 9944f2e commit d3bb902

7 files changed

Lines changed: 109 additions & 16 deletions

File tree

‎integ-test/src/test/java/org/opensearch/sql/ppl/PrometheusCatalogCommandsIT.java‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -46,12 +46,12 @@ public void testSourceMetricCommand() {
4646
@SneakyThrows
4747
public void testMetricAvgAggregationCommand() {
4848
JSONObject response =
49-
executeQuery("source=my_prometheus.prometheus_http_requests_total | stats avg(@value) by span(@timestamp, 15s), handler, job");
49+
executeQuery("source=`my_prometheus`.`prometheus_http_requests_total` | stats avg(@value) as `agg` by span(@timestamp, 15s), `handler`, `job`");
5050
verifySchema(response,
51-
schema("avg(@value)", "double"),
51+
schema("agg", "double"),
5252
schema("span(@timestamp,15s)", "timestamp"),
53-
schema("handler", "string"),
54-
schema("job", "string"));
53+
schema("`handler`", "string"),
54+
schema("`job`", "string"));
5555
Assertions.assertTrue(response.getInt("size") > 0);
5656
Assertions.assertEquals(4, response.getJSONArray("datarows").getJSONArray(0).length());
5757
JSONArray firstRow = response.getJSONArray("datarows").getJSONArray(0);
@@ -65,11 +65,11 @@ public void testMetricAvgAggregationCommand() {
6565
@SneakyThrows
6666
public void testMetricAvgAggregationCommandWithAlias() {
6767
JSONObject response =
68-
executeQuery("source=my_prometheus.prometheus_http_requests_total | stats avg(@value) as agg by span(@timestamp, 15s), handler, job");
68+
executeQuery("source=my_prometheus.prometheus_http_requests_total | stats avg(@value) as agg by span(@timestamp, 15s), `handler`, job");
6969
verifySchema(response,
7070
schema("agg", "double"),
7171
schema("span(@timestamp,15s)", "timestamp"),
72-
schema("handler", "string"),
72+
schema("`handler`", "string"),
7373
schema("job", "string"));
7474
Assertions.assertTrue(response.getInt("size") > 0);
7575
Assertions.assertEquals(4, response.getJSONArray("datarows").getJSONArray(0).length());

‎prometheus/src/main/java/org/opensearch/sql/prometheus/response/PrometheusResponse.java‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55

66
package org.opensearch.sql.prometheus.response;
77

8-
import static org.opensearch.sql.data.type.ExprCoreType.DOUBLE;
98
import static org.opensearch.sql.data.type.ExprCoreType.INTEGER;
109
import static org.opensearch.sql.data.type.ExprCoreType.LONG;
1110
import static org.opensearch.sql.prometheus.data.constants.PrometheusFieldConstants.LABELS;
@@ -16,6 +15,7 @@
1615
import java.util.LinkedHashMap;
1716
import java.util.List;
1817
import lombok.NonNull;
18+
import org.apache.commons.lang3.StringUtils;
1919
import org.json.JSONArray;
2020
import org.json.JSONObject;
2121
import org.opensearch.sql.data.model.ExprDoubleValue;
@@ -26,6 +26,8 @@
2626
import org.opensearch.sql.data.model.ExprTupleValue;
2727
import org.opensearch.sql.data.model.ExprValue;
2828
import org.opensearch.sql.data.type.ExprType;
29+
import org.opensearch.sql.expression.NamedExpression;
30+
import org.opensearch.sql.expression.ReferenceExpression;
2931
import org.opensearch.sql.prometheus.storage.model.PrometheusResponseFieldNames;
3032

3133
public class PrometheusResponse implements Iterable<ExprValue> {
@@ -100,7 +102,7 @@ public Iterator<ExprValue> iterator() {
100102

101103
private void insertLabels(LinkedHashMap<String, ExprValue> linkedHashMap, JSONObject metric) {
102104
for (String key : metric.keySet()) {
103-
linkedHashMap.put(key, new ExprStringValue(metric.getString(key)));
105+
linkedHashMap.put(getKey(key), new ExprStringValue(metric.getString(key)));
104106
}
105107
}
106108

@@ -113,4 +115,18 @@ private ExprValue getValue(JSONArray jsonArray, Integer index, ExprType exprType
113115
return new ExprDoubleValue(jsonArray.getDouble(index));
114116
}
115117

118+
private String getKey(String key) {
119+
if (this.prometheusResponseFieldNames.getGroupByList() == null) {
120+
return key;
121+
} else {
122+
return this.prometheusResponseFieldNames.getGroupByList().stream()
123+
.filter(expression -> expression.getDelegated() instanceof ReferenceExpression)
124+
.filter(expression
125+
-> ((ReferenceExpression) expression.getDelegated()).getAttr().equals(key))
126+
.findFirst()
127+
.map(NamedExpression::getName)
128+
.orElse(key);
129+
}
130+
}
131+
116132
}

‎prometheus/src/main/java/org/opensearch/sql/prometheus/storage/implementor/PrometheusDefaultImplementor.java‎

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,25 +7,18 @@
77

88
package org.opensearch.sql.prometheus.storage.implementor;
99

10-
import static org.opensearch.sql.data.type.ExprCoreType.STRING;
11-
import static org.opensearch.sql.prometheus.data.constants.PrometheusFieldConstants.LABELS;
12-
13-
import java.util.ArrayList;
1410
import java.util.List;
1511
import java.util.Optional;
1612
import lombok.RequiredArgsConstructor;
1713
import org.apache.commons.math3.util.Pair;
1814
import org.opensearch.sql.common.utils.StringUtils;
1915
import org.opensearch.sql.expression.Expression;
2016
import org.opensearch.sql.expression.NamedExpression;
21-
import org.opensearch.sql.expression.ReferenceExpression;
2217
import org.opensearch.sql.expression.span.SpanExpression;
2318
import org.opensearch.sql.planner.DefaultImplementor;
2419
import org.opensearch.sql.planner.logical.LogicalPlan;
25-
import org.opensearch.sql.planner.logical.LogicalProject;
2620
import org.opensearch.sql.planner.logical.LogicalRelation;
2721
import org.opensearch.sql.planner.physical.PhysicalPlan;
28-
import org.opensearch.sql.planner.physical.ProjectOperator;
2922
import org.opensearch.sql.prometheus.planner.logical.PrometheusLogicalMetricAgg;
3023
import org.opensearch.sql.prometheus.planner.logical.PrometheusLogicalMetricScan;
3124
import org.opensearch.sql.prometheus.storage.PrometheusMetricScan;
@@ -130,6 +123,7 @@ private void setPrometheusResponseFieldNames(PrometheusLogicalMetricAgg node,
130123
prometheusResponseFieldNames.setValueFieldName(node.getAggregatorList().get(0).getName());
131124
prometheusResponseFieldNames.setValueType(node.getAggregatorList().get(0).type());
132125
prometheusResponseFieldNames.setTimestampFieldName(spanExpression.get().getNameOrAlias());
126+
prometheusResponseFieldNames.setGroupByList(node.getGroupByList());
133127
context.setPrometheusResponseFieldNames(prometheusResponseFieldNames);
134128
}
135129

‎prometheus/src/main/java/org/opensearch/sql/prometheus/storage/model/PrometheusResponseFieldNames.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,11 @@
1111
import static org.opensearch.sql.prometheus.data.constants.PrometheusFieldConstants.TIMESTAMP;
1212
import static org.opensearch.sql.prometheus.data.constants.PrometheusFieldConstants.VALUE;
1313

14+
import java.util.List;
1415
import lombok.Getter;
1516
import lombok.Setter;
1617
import org.opensearch.sql.data.type.ExprType;
18+
import org.opensearch.sql.expression.NamedExpression;
1719

1820

1921
@Getter
@@ -23,5 +25,6 @@ public class PrometheusResponseFieldNames {
2325
private String valueFieldName = VALUE;
2426
private ExprType valueType = DOUBLE;
2527
private String timestampFieldName = TIMESTAMP;
28+
private List<NamedExpression> groupByList;
2629

2730
}

‎prometheus/src/main/java/org/opensearch/sql/prometheus/storage/querybuilder/AggregationQueryBuilder.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,14 @@
77

88
package org.opensearch.sql.prometheus.storage.querybuilder;
99

10+
import java.sql.Ref;
1011
import java.util.List;
1112
import java.util.Set;
1213
import java.util.stream.Collectors;
1314
import lombok.NoArgsConstructor;
15+
import org.apache.commons.lang3.StringUtils;
1416
import org.opensearch.sql.expression.NamedExpression;
17+
import org.opensearch.sql.expression.ReferenceExpression;
1518
import org.opensearch.sql.expression.aggregation.NamedAggregator;
1619
import org.opensearch.sql.expression.function.BuiltinFunctionName;
1720
import org.opensearch.sql.expression.span.SpanExpression;
@@ -63,7 +66,10 @@ public static String build(List<NamedAggregator> namedAggregatorList,
6366
if (groupByList.size() > 0) {
6467
aggregateQuery.append("by(");
6568
aggregateQuery.append(
66-
groupByList.stream().map(NamedExpression::getName).collect(Collectors.joining(", ")));
69+
groupByList.stream()
70+
.filter(expression -> expression.getDelegated() instanceof ReferenceExpression)
71+
.map(expression -> ((ReferenceExpression) expression.getDelegated()).getAttr())
72+
.collect(Collectors.joining(", ")));
6773
aggregateQuery.append(")");
6874
}
6975
}

‎prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusMetricScanTest.java‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import static org.mockito.Mockito.when;
1212
import static org.opensearch.sql.data.type.ExprCoreType.INTEGER;
1313
import static org.opensearch.sql.data.type.ExprCoreType.LONG;
14+
import static org.opensearch.sql.data.type.ExprCoreType.STRING;
1415
import static org.opensearch.sql.prometheus.constants.TestConstants.ENDTIME;
1516
import static org.opensearch.sql.prometheus.constants.TestConstants.QUERY;
1617
import static org.opensearch.sql.prometheus.constants.TestConstants.STARTTIME;
@@ -22,6 +23,7 @@
2223

2324
import java.io.IOException;
2425
import java.time.Instant;
26+
import java.util.Collections;
2527
import java.util.LinkedHashMap;
2628
import lombok.SneakyThrows;
2729
import org.json.JSONObject;
@@ -36,6 +38,7 @@
3638
import org.opensearch.sql.data.model.ExprStringValue;
3739
import org.opensearch.sql.data.model.ExprTimestampValue;
3840
import org.opensearch.sql.data.model.ExprTupleValue;
41+
import org.opensearch.sql.expression.DSL;
3942
import org.opensearch.sql.prometheus.client.PrometheusClient;
4043
import org.opensearch.sql.prometheus.storage.model.PrometheusResponseFieldNames;
4144

@@ -163,6 +166,49 @@ void testQueryResponseIteratorWithGivenPrometheusResponseWithLongInAggType() {
163166
Assertions.assertFalse(prometheusMetricScan.hasNext());
164167
}
165168

169+
@Test
170+
@SneakyThrows
171+
void testQueryResponseIteratorWithGivenPrometheusResponseWithBackQuotedFieldNames() {
172+
PrometheusResponseFieldNames prometheusResponseFieldNames
173+
= new PrometheusResponseFieldNames();
174+
prometheusResponseFieldNames.setValueFieldName("testAgg");
175+
prometheusResponseFieldNames.setValueType(LONG);
176+
prometheusResponseFieldNames.setTimestampFieldName(TIMESTAMP);
177+
prometheusResponseFieldNames.setGroupByList(
178+
Collections.singletonList(DSL.named("`instance`", DSL.ref("instance", STRING))));
179+
PrometheusMetricScan prometheusMetricScan = new PrometheusMetricScan(prometheusClient);
180+
prometheusMetricScan.setPrometheusResponseFieldNames(prometheusResponseFieldNames);
181+
prometheusMetricScan.getRequest().setPromQl(QUERY);
182+
prometheusMetricScan.getRequest().setStartTime(STARTTIME);
183+
prometheusMetricScan.getRequest().setEndTime(ENDTIME);
184+
prometheusMetricScan.getRequest().setStep(STEP);
185+
186+
when(prometheusClient.queryRange(any(), any(), any(), any()))
187+
.thenReturn(new JSONObject(getJson("query_range_result.json")));
188+
prometheusMetricScan.open();
189+
Assertions.assertTrue(prometheusMetricScan.hasNext());
190+
ExprTupleValue firstRow = new ExprTupleValue(new LinkedHashMap<>() {{
191+
put(TIMESTAMP, new ExprTimestampValue(Instant.ofEpochMilli(1435781430781L)));
192+
put("testAgg", new ExprLongValue(1));
193+
put("`instance`", new ExprStringValue("localhost:9090"));
194+
put("__name__", new ExprStringValue("up"));
195+
put("job", new ExprStringValue("prometheus"));
196+
}
197+
});
198+
assertEquals(firstRow, prometheusMetricScan.next());
199+
Assertions.assertTrue(prometheusMetricScan.hasNext());
200+
ExprTupleValue secondRow = new ExprTupleValue(new LinkedHashMap<>() {{
201+
put(TIMESTAMP, new ExprTimestampValue(Instant.ofEpochMilli(1435781430781L)));
202+
put("testAgg", new ExprLongValue(0));
203+
put("`instance`", new ExprStringValue("localhost:9091"));
204+
put("__name__", new ExprStringValue("up"));
205+
put("job", new ExprStringValue("node"));
206+
}
207+
});
208+
assertEquals(secondRow, prometheusMetricScan.next());
209+
Assertions.assertFalse(prometheusMetricScan.hasNext());
210+
}
211+
166212
@Test
167213
@SneakyThrows
168214
void testQueryResponseIteratorForQueryRangeFunction() {
@@ -247,6 +293,7 @@ void testEmptyQueryWithException() {
247293
runtimeException.getMessage());
248294
}
249295

296+
250297
@Test
251298
@SneakyThrows
252299
void testExplain() {

‎prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusMetricTableTest.java‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -739,4 +739,31 @@ void testOptimize() {
739739
assertEquals(inputPlan, optimizedPlan);
740740
}
741741

742+
@Test
743+
void testImplementPrometheusQueryWithBackQuotedFieldNamesInStatsQuery() {
744+
745+
PrometheusMetricTable prometheusMetricTable =
746+
new PrometheusMetricTable(client, "prometheus_http_total_requests");
747+
748+
749+
// IndexScanAgg with Filter
750+
PhysicalPlan plan = prometheusMetricTable.implement(
751+
indexScanAgg("prometheus_http_total_requests",
752+
dsl.and(dsl.equal(DSL.ref("code", STRING), DSL.literal(stringValue("200"))),
753+
dsl.equal(DSL.ref("handler", STRING), DSL.literal(stringValue("/ready/")))),
754+
ImmutableList
755+
.of(named("AVG(@value)",
756+
dsl.avg(DSL.ref("@value", INTEGER)))),
757+
ImmutableList.of(named("`job`", DSL.ref("job", STRING)),
758+
named("span", DSL.span(DSL.ref("@timestamp", ExprCoreType.TIMESTAMP),
759+
DSL.literal(40), "s")))));
760+
assertTrue(plan instanceof PrometheusMetricScan);
761+
PrometheusQueryRequest prometheusQueryRequest = ((PrometheusMetricScan) plan).getRequest();
762+
assertEquals(
763+
"avg by(job) (avg_over_time"
764+
+ "(prometheus_http_total_requests{code=\"200\" , handler=\"/ready/\"}[40s]))",
765+
prometheusQueryRequest.getPromQl());
766+
767+
}
768+
742769
}

0 commit comments

Comments
 (0)