Skip to content

Commit a528031

Browse files
Enable Table Function and PromQL function (#1719) (#1744)
1 parent 7fe2fb4 commit a528031

File tree

16 files changed

+598
-15
lines changed

16 files changed

+598
-15
lines changed

docs/user/ppl/admin/prometheus_connector.rst

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -186,3 +186,27 @@ Example queries
186186
| 11 | "2022-11-03 07:18:64" | "/-/metrics" | 500 |
187187
+------------+------------------------+--------------------------------+---------------+
188188

189+
PromQL Support for prometheus Connector
190+
==========================================
191+
192+
`query_range` Table Function
193+
----------------------------
194+
Prometheus connector offers `query_range` table function. This table function can be used to query metrics in a specific time range using promQL.
195+
The function takes inputs similar to parameters mentioned for query range api mentioned here: https://prometheus.io/docs/prometheus/latest/querying/api/
196+
Arguments should be either passed by name or positionArguments should be either passed by name or position.
197+
`source=my_prometheus.query_range('prometheus_http_requests_total', 1686694425, 1686700130, 14)`
198+
or
199+
`source=my_prometheus.query_range(query='prometheus_http_requests_total', starttime=1686694425, endtime=1686700130, step=14)`
200+
Example::
201+
202+
> source=my_prometheus.query_range('prometheus_http_requests_total', 1686694425, 1686700130, 14)
203+
+------------+------------------------+--------------------------------+---------------+-------------+-------------+
204+
| @value | @timestamp | handler | code | instance | job |
205+
|------------+------------------------+--------------------------------+---------------+-------------+-------------|
206+
| 5 | "2022-11-03 07:18:14" | "/-/ready" | 200 | 192.15.1.1 | prometheus |
207+
| 3 | "2022-11-03 07:18:24" | "/-/ready" | 200 | 192.15.1.1 | prometheus |
208+
| 7 | "2022-11-03 07:18:34" | "/-/ready" | 200 | 192.15.1.1 | prometheus |
209+
| 2 | "2022-11-03 07:18:44" | "/-/ready" | 400 | 192.15.2.1 | prometheus |
210+
| 9 | "2022-11-03 07:18:54" | "/-/promql" | 400 | 192.15.2.1 | prometheus |
211+
| 11 | "2022-11-03 07:18:64" |"/-/metrics" | 500 | 192.15.2.1 | prometheus |
212+
+------------+------------------------+--------------------------------+---------------+-------------+-------------+

ppl/src/main/antlr/OpenSearchPPLParser.g4

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -183,8 +183,11 @@ mlArg
183183
fromClause
184184
: SOURCE EQUAL tableSourceClause
185185
| INDEX EQUAL tableSourceClause
186+
| SOURCE EQUAL tableFunction
187+
| INDEX EQUAL tableFunction
186188
;
187189

190+
188191
tableSourceClause
189192
: tableSource (COMMA tableSource)*
190193
;

ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import static org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParser.SearchFromFilterContext;
2121
import static org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParser.SortCommandContext;
2222
import static org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParser.StatsCommandContext;
23+
import static org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParser.TableFunctionContext;
2324
import static org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParser.TableSourceClauseContext;
2425
import static org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParser.TopCommandContext;
2526
import static org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParser.WhereCommandContext;
@@ -33,7 +34,6 @@
3334
import java.util.List;
3435
import java.util.Optional;
3536
import java.util.stream.Collectors;
36-
import lombok.Generated;
3737
import lombok.RequiredArgsConstructor;
3838
import org.antlr.v4.runtime.ParserRuleContext;
3939
import org.antlr.v4.runtime.Token;
@@ -46,6 +46,7 @@
4646
import org.opensearch.sql.ast.expression.Map;
4747
import org.opensearch.sql.ast.expression.ParseMethod;
4848
import org.opensearch.sql.ast.expression.QualifiedName;
49+
import org.opensearch.sql.ast.expression.UnresolvedArgument;
4950
import org.opensearch.sql.ast.expression.UnresolvedExpression;
5051
import org.opensearch.sql.ast.tree.AD;
5152
import org.opensearch.sql.ast.tree.Aggregation;
@@ -62,6 +63,7 @@
6263
import org.opensearch.sql.ast.tree.Relation;
6364
import org.opensearch.sql.ast.tree.Rename;
6465
import org.opensearch.sql.ast.tree.Sort;
66+
import org.opensearch.sql.ast.tree.TableFunction;
6567
import org.opensearch.sql.ast.tree.UnresolvedPlan;
6668
import org.opensearch.sql.common.utils.StringUtils;
6769
import org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParser;
@@ -346,7 +348,11 @@ public UnresolvedPlan visitTopCommand(TopCommandContext ctx) {
346348
*/
347349
@Override
348350
public UnresolvedPlan visitFromClause(FromClauseContext ctx) {
349-
return visitTableSourceClause(ctx.tableSourceClause());
351+
if (ctx.tableFunction() != null) {
352+
return visitTableFunction(ctx.tableFunction());
353+
} else {
354+
return visitTableSourceClause(ctx.tableSourceClause());
355+
}
350356
}
351357

352358
@Override
@@ -357,10 +363,16 @@ public UnresolvedPlan visitTableSourceClause(TableSourceClauseContext ctx) {
357363
}
358364

359365
@Override
360-
@Generated //To exclude from jacoco..will remove https://github.com/opensearch-project/sql/issues/1019
361-
public UnresolvedPlan visitTableFunction(OpenSearchPPLParser.TableFunctionContext ctx) {
362-
//<TODO>
363-
return null;
366+
public UnresolvedPlan visitTableFunction(TableFunctionContext ctx) {
367+
ImmutableList.Builder<UnresolvedExpression> builder = ImmutableList.builder();
368+
ctx.functionArgs().functionArg().forEach(arg
369+
-> {
370+
String argName = (arg.ident() != null) ? arg.ident().getText() : null;
371+
builder.add(
372+
new UnresolvedArgument(argName,
373+
this.internalVisitExpression(arg.valueExpression())));
374+
});
375+
return new TableFunction(this.internalVisitExpression(ctx.qualifiedName()), builder.build());
364376
}
365377

366378
/**

ppl/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66

77
package org.opensearch.sql.ppl.parser;
88

9-
import static org.opensearch.sql.ast.dsl.AstDSL.qualifiedName;
109
import static org.opensearch.sql.expression.function.BuiltinFunctionName.IS_NOT_NULL;
1110
import static org.opensearch.sql.expression.function.BuiltinFunctionName.IS_NULL;
1211
import static org.opensearch.sql.expression.function.BuiltinFunctionName.POSITION;

ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryDataAnonymizer.java

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -101,10 +101,13 @@ public String visitRelation(Relation node, String context) {
101101
}
102102

103103
@Override
104-
@Generated //To exclude from jacoco..will remove https://github.com/opensearch-project/sql/issues/1019
105104
public String visitTableFunction(TableFunction node, String context) {
106-
//<TODO>
107-
return null;
105+
String arguments =
106+
node.getArguments().stream()
107+
.map(unresolvedExpression
108+
-> this.expressionAnalyzer.analyze(unresolvedExpression, context))
109+
.collect(Collectors.joining(","));
110+
return StringUtils.format("source=%s(%s)", node.getFunctionName().toString(), arguments);
108111
}
109112

110113
@Override

ppl/src/test/java/org/opensearch/sql/ppl/parser/AstBuilderTest.java

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,6 @@ public void testSearchCommandWithDotInIndexName() {
112112
);
113113
}
114114

115-
@Ignore
116115
@Test
117116
public void testSearchWithPrometheusQueryRangeWithPositionedArguments() {
118117
assertEqual("search source = prometheus.query_range(\"test{code='200'}\",1234, 12345, 3)",
@@ -124,7 +123,6 @@ public void testSearchWithPrometheusQueryRangeWithPositionedArguments() {
124123
));
125124
}
126125

127-
@Ignore
128126
@Test
129127
public void testSearchWithPrometheusQueryRangeWithNamedArguments() {
130128
assertEqual("search source = prometheus.query_range(query = \"test{code='200'}\", "

ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryDataAnonymizerTest.java

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
import static org.opensearch.sql.ast.dsl.AstDSL.relation;
1313

1414
import java.util.Collections;
15-
import org.junit.Ignore;
1615
import org.junit.Test;
1716
import org.junit.runner.RunWith;
1817
import org.mockito.junit.MockitoJUnitRunner;
@@ -36,7 +35,6 @@ public void testSearchCommand() {
3635
}
3736

3837
@Test
39-
@Ignore
4038
public void testTableFunctionCommand() {
4139
assertEquals("source=prometheus.query_range(***,***,***,***)",
4240
anonymize("source=prometheus.query_range('afsd',123,123,3)")
Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.prometheus.functions.response;
7+
8+
import static org.opensearch.sql.prometheus.data.constants.PrometheusFieldConstants.VALUE;
9+
10+
import java.time.Instant;
11+
import java.util.ArrayList;
12+
import java.util.Iterator;
13+
import java.util.LinkedHashMap;
14+
import java.util.List;
15+
import org.jetbrains.annotations.NotNull;
16+
import org.json.JSONArray;
17+
import org.json.JSONObject;
18+
import org.opensearch.sql.data.model.ExprDoubleValue;
19+
import org.opensearch.sql.data.model.ExprStringValue;
20+
import org.opensearch.sql.data.model.ExprTimestampValue;
21+
import org.opensearch.sql.data.model.ExprTupleValue;
22+
import org.opensearch.sql.data.model.ExprValue;
23+
import org.opensearch.sql.data.type.ExprCoreType;
24+
import org.opensearch.sql.executor.ExecutionEngine;
25+
import org.opensearch.sql.prometheus.data.constants.PrometheusFieldConstants;
26+
27+
/**
28+
* Default implementation of QueryRangeFunctionResponseHandle.
29+
*/
30+
public class DefaultQueryRangeFunctionResponseHandle implements QueryRangeFunctionResponseHandle {
31+
32+
private final JSONObject responseObject;
33+
private Iterator<ExprValue> responseIterator;
34+
private ExecutionEngine.Schema schema;
35+
36+
/**
37+
* Constructor.
38+
*
39+
* @param responseObject Prometheus responseObject.
40+
*/
41+
public DefaultQueryRangeFunctionResponseHandle(JSONObject responseObject) {
42+
this.responseObject = responseObject;
43+
constructIteratorAndSchema();
44+
}
45+
46+
private void constructIteratorAndSchema() {
47+
List<ExprValue> result = new ArrayList<>();
48+
List<ExecutionEngine.Schema.Column> columnList = new ArrayList<>();
49+
if ("matrix".equals(responseObject.getString("resultType"))) {
50+
JSONArray itemArray = responseObject.getJSONArray("result");
51+
for (int i = 0; i < itemArray.length(); i++) {
52+
JSONObject item = itemArray.getJSONObject(i);
53+
JSONObject metric = item.getJSONObject("metric");
54+
JSONArray values = item.getJSONArray("values");
55+
if (i == 0) {
56+
columnList = getColumnList(metric);
57+
}
58+
for (int j = 0; j < values.length(); j++) {
59+
LinkedHashMap<String, ExprValue> linkedHashMap =
60+
extractRow(metric, values.getJSONArray(j), columnList);
61+
result.add(new ExprTupleValue(linkedHashMap));
62+
}
63+
}
64+
} else {
65+
throw new RuntimeException(String.format("Unexpected Result Type: %s during Prometheus "
66+
+ "Response Parsing. 'matrix' resultType is expected",
67+
responseObject.getString("resultType")));
68+
}
69+
this.schema = new ExecutionEngine.Schema(columnList);
70+
this.responseIterator = result.iterator();
71+
}
72+
73+
@NotNull
74+
private static LinkedHashMap<String, ExprValue> extractRow(JSONObject metric,
75+
JSONArray values, List<ExecutionEngine.Schema.Column> columnList) {
76+
LinkedHashMap<String, ExprValue> linkedHashMap = new LinkedHashMap<>();
77+
for (ExecutionEngine.Schema.Column column : columnList) {
78+
if (PrometheusFieldConstants.TIMESTAMP.equals(column.getName())) {
79+
linkedHashMap.put(PrometheusFieldConstants.TIMESTAMP,
80+
new ExprTimestampValue(Instant.ofEpochMilli((long) (values.getDouble(0) * 1000))));
81+
} else if (column.getName().equals(VALUE)) {
82+
linkedHashMap.put(VALUE, new ExprDoubleValue(values.getDouble(1)));
83+
} else {
84+
linkedHashMap.put(column.getName(),
85+
new ExprStringValue(metric.getString(column.getName())));
86+
}
87+
}
88+
return linkedHashMap;
89+
}
90+
91+
92+
private List<ExecutionEngine.Schema.Column> getColumnList(JSONObject metric) {
93+
List<ExecutionEngine.Schema.Column> columnList = new ArrayList<>();
94+
columnList.add(new ExecutionEngine.Schema.Column(PrometheusFieldConstants.TIMESTAMP,
95+
PrometheusFieldConstants.TIMESTAMP, ExprCoreType.TIMESTAMP));
96+
columnList.add(new ExecutionEngine.Schema.Column(VALUE, VALUE, ExprCoreType.DOUBLE));
97+
for (String key : metric.keySet()) {
98+
columnList.add(new ExecutionEngine.Schema.Column(key, key, ExprCoreType.STRING));
99+
}
100+
return columnList;
101+
}
102+
103+
@Override
104+
public boolean hasNext() {
105+
return responseIterator.hasNext();
106+
}
107+
108+
@Override
109+
public ExprValue next() {
110+
return responseIterator.next();
111+
}
112+
113+
@Override
114+
public ExecutionEngine.Schema schema() {
115+
return schema;
116+
}
117+
}
Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.prometheus.functions.response;
7+
8+
import org.opensearch.sql.data.model.ExprValue;
9+
import org.opensearch.sql.executor.ExecutionEngine;
10+
11+
/**
12+
* Handle Prometheus response.
13+
*/
14+
public interface QueryRangeFunctionResponseHandle {
15+
16+
/**
17+
* Return true if Prometheus response has more result.
18+
*/
19+
boolean hasNext();
20+
21+
/**
22+
* Return Prometheus response as {@link ExprValue}. Attention, the method must been called when
23+
* hasNext return true.
24+
*/
25+
ExprValue next();
26+
27+
/**
28+
* Return ExecutionEngine.Schema of the Prometheus response.
29+
*/
30+
ExecutionEngine.Schema schema();
31+
}
Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
/*
2+
*
3+
* * Copyright OpenSearch Contributors
4+
* * SPDX-License-Identifier: Apache-2.0
5+
*
6+
*/
7+
8+
package org.opensearch.sql.prometheus.functions.scan;
9+
10+
import lombok.AllArgsConstructor;
11+
import org.opensearch.sql.planner.logical.LogicalProject;
12+
import org.opensearch.sql.prometheus.client.PrometheusClient;
13+
import org.opensearch.sql.prometheus.request.PrometheusQueryRequest;
14+
import org.opensearch.sql.storage.TableScanOperator;
15+
import org.opensearch.sql.storage.read.TableScanBuilder;
16+
17+
/**
18+
* TableScanBuilder for query_range table function of prometheus connector.
19+
* we can merge this when we refactor for existing
20+
* ppl queries based on prometheus connector.
21+
*/
22+
@AllArgsConstructor
23+
public class QueryRangeFunctionTableScanBuilder extends TableScanBuilder {
24+
25+
private final PrometheusClient prometheusClient;
26+
27+
private final PrometheusQueryRequest prometheusQueryRequest;
28+
29+
@Override
30+
public TableScanOperator build() {
31+
return new QueryRangeFunctionTableScanOperator(prometheusClient, prometheusQueryRequest);
32+
}
33+
34+
@Override
35+
public boolean pushDownProject(LogicalProject project) {
36+
return true;
37+
}
38+
}

0 commit comments

Comments
 (0)