query/
parser.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::any::Any;
16use std::collections::HashMap;
17use std::sync::Arc;
18use std::time::{Duration, SystemTime};
19
20use chrono::DateTime;
21use common_error::ext::{BoxedError, PlainError};
22use common_error::status_code::StatusCode;
23use common_telemetry::tracing;
24use promql_parser::parser::Expr::Extension;
25use promql_parser::parser::ast::{Extension as NodeExtension, ExtensionExpr};
26use promql_parser::parser::value::ValueType;
27use promql_parser::parser::{EvalStmt, Expr};
28use session::context::QueryContextRef;
29use snafu::{OptionExt, ResultExt};
30use sql::dialect::GreptimeDbDialect;
31use sql::parser::{ParseOptions, ParserContext};
32use sql::statements::statement::Statement;
33
34use crate::error::{
35    AddSystemTimeOverflowSnafu, MultipleStatementsSnafu, ParseFloatSnafu, ParseTimestampSnafu,
36    QueryParseSnafu, Result, TryIntoDurationSnafu, UnimplementedSnafu,
37};
38use crate::metrics::{PARSE_PROMQL_ELAPSED, PARSE_SQL_ELAPSED};
39
40pub const DEFAULT_LOOKBACK_STRING: &str = "5m";
41pub const EXPLAIN_NODE_NAME: &str = "EXPLAIN";
42pub const EXPLAIN_VERBOSE_NODE_NAME: &str = "EXPLAIN VERBOSE";
43pub const ANALYZE_NODE_NAME: &str = "ANALYZE";
44pub const ANALYZE_VERBOSE_NODE_NAME: &str = "ANALYZE VERBOSE";
45pub const ALIAS_NODE_NAME: &str = "ALIAS";
46
47#[derive(Debug, Clone)]
48pub enum QueryStatement {
49    Sql(Statement),
50    // The optional String is the alias name
51    Promql(EvalStmt, Option<String>),
52}
53
54impl QueryStatement {
55    pub fn post_process(self, params: HashMap<String, String>) -> Result<QueryStatement> {
56        match self {
57            QueryStatement::Sql(_) => UnimplementedSnafu {
58                operation: "sql post process",
59            }
60            .fail(),
61            QueryStatement::Promql(mut eval_stmt, alias) => {
62                let node_name = match params.get("name") {
63                    Some(name) => name.as_str(),
64                    None => "",
65                };
66                let extension_node = Self::create_extension_node(node_name, eval_stmt.expr);
67                eval_stmt.expr = Extension(extension_node.unwrap());
68                Ok(QueryStatement::Promql(eval_stmt, alias))
69            }
70        }
71    }
72
73    fn create_extension_node(node_name: &str, expr: Expr) -> Option<NodeExtension> {
74        match node_name {
75            ANALYZE_NODE_NAME => Some(NodeExtension {
76                expr: Arc::new(AnalyzeExpr { expr }),
77            }),
78            ANALYZE_VERBOSE_NODE_NAME => Some(NodeExtension {
79                expr: Arc::new(AnalyzeVerboseExpr { expr }),
80            }),
81            EXPLAIN_NODE_NAME => Some(NodeExtension {
82                expr: Arc::new(ExplainExpr { expr }),
83            }),
84            EXPLAIN_VERBOSE_NODE_NAME => Some(NodeExtension {
85                expr: Arc::new(ExplainVerboseExpr { expr }),
86            }),
87            _ => None,
88        }
89    }
90}
91
92#[derive(Debug, Clone, PartialEq, Eq)]
93pub struct PromQuery {
94    pub query: String,
95    pub start: String,
96    pub end: String,
97    pub step: String,
98    pub lookback: String,
99    pub alias: Option<String>,
100}
101
102impl Default for PromQuery {
103    fn default() -> Self {
104        PromQuery {
105            query: String::new(),
106            start: String::from("0"),
107            end: String::from("0"),
108            step: String::from("5m"),
109            lookback: String::from(DEFAULT_LOOKBACK_STRING),
110            alias: None,
111        }
112    }
113}
114
115/// Query language parser, supports parsing SQL and PromQL
116pub struct QueryLanguageParser {}
117
118impl QueryLanguageParser {
119    /// Try to parse SQL with GreptimeDB dialect, return the statement when success.
120    pub fn parse_sql(sql: &str, query_ctx: &QueryContextRef) -> Result<QueryStatement> {
121        let _timer = PARSE_SQL_ELAPSED.start_timer();
122        let scheduled_time =
123            crate::options::parse_scheduled_time_datetime(&query_ctx.extensions())?;
124        let mut statement = ParserContext::create_with_dialect(
125            sql,
126            &GreptimeDbDialect {},
127            ParseOptions { scheduled_time },
128        )
129        .map_err(BoxedError::new)
130        .context(QueryParseSnafu { query: sql })?;
131        if statement.len() != 1 {
132            MultipleStatementsSnafu {
133                query: sql.to_string(),
134            }
135            .fail()
136        } else {
137            Ok(QueryStatement::Sql(statement.pop().unwrap()))
138        }
139    }
140
141    /// Try to parse PromQL, return the statement when success.
142    #[tracing::instrument(skip_all)]
143    pub fn parse_promql(query: &PromQuery, _query_ctx: &QueryContextRef) -> Result<QueryStatement> {
144        let _timer = PARSE_PROMQL_ELAPSED.start_timer();
145
146        let expr = promql_parser::parser::parse(&query.query)
147            .map_err(|msg| BoxedError::new(PlainError::new(msg, StatusCode::InvalidArguments)))
148            .context(QueryParseSnafu {
149                query: &query.query,
150            })?;
151
152        let start = Self::parse_promql_timestamp(&query.start)
153            .map_err(BoxedError::new)
154            .context(QueryParseSnafu {
155                query: &query.query,
156            })?;
157
158        let end = Self::parse_promql_timestamp(&query.end)
159            .map_err(BoxedError::new)
160            .context(QueryParseSnafu {
161                query: &query.query,
162            })?;
163
164        let step = query
165            .step
166            .parse::<u64>()
167            .map(Duration::from_secs)
168            .or_else(|_| promql_parser::util::parse_duration(&query.step))
169            .map_err(|msg| BoxedError::new(PlainError::new(msg, StatusCode::InvalidArguments)))
170            .context(QueryParseSnafu {
171                query: &query.query,
172            })?;
173
174        let lookback_delta = query
175            .lookback
176            .parse::<u64>()
177            .map(Duration::from_secs)
178            .or_else(|_| promql_parser::util::parse_duration(&query.lookback))
179            .map_err(|msg| BoxedError::new(PlainError::new(msg, StatusCode::InvalidArguments)))
180            .context(QueryParseSnafu {
181                query: &query.query,
182            })?;
183
184        let eval_stmt = EvalStmt {
185            expr,
186            start,
187            end,
188            interval: step,
189            lookback_delta,
190        };
191        if let Some(alias) = &query.alias {
192            let eval_stmt = Self::apply_alias_extension(eval_stmt, alias);
193            return Ok(QueryStatement::Promql(eval_stmt, query.alias.clone()));
194        }
195        Ok(QueryStatement::Promql(eval_stmt, None))
196    }
197
198    pub(crate) fn apply_alias_extension(mut eval_stmt: EvalStmt, alias: &str) -> EvalStmt {
199        let expr = eval_stmt.expr;
200        eval_stmt.expr = Extension(NodeExtension {
201            expr: Arc::new(AliasExpr {
202                expr,
203                alias: alias.to_string(),
204            }),
205        });
206        eval_stmt
207    }
208
209    pub fn parse_promql_timestamp(timestamp: &str) -> Result<SystemTime> {
210        // try rfc3339 format
211        let rfc3339_result = DateTime::parse_from_rfc3339(timestamp)
212            .context(ParseTimestampSnafu { raw: timestamp })
213            .map(Into::<SystemTime>::into);
214
215        // shorthand
216        if rfc3339_result.is_ok() {
217            return rfc3339_result;
218        }
219
220        // try float format
221        let secs = timestamp
222            .parse::<f64>()
223            .context(ParseFloatSnafu { raw: timestamp })
224            // also report rfc3339 error if float parsing fails
225            .map_err(|_| rfc3339_result.unwrap_err())?;
226
227        let duration =
228            Duration::try_from_secs_f64(secs).context(TryIntoDurationSnafu { raw: timestamp })?;
229        SystemTime::UNIX_EPOCH
230            .checked_add(duration)
231            .context(AddSystemTimeOverflowSnafu { duration })
232    }
233}
234
235macro_rules! define_node_ast_extension {
236    ($name:ident, $name_expr:ident, $expr_type:ty, $extension_name:expr) => {
237        /// The implementation of the `$name_expr` extension AST node
238        #[derive(Debug, Clone)]
239        pub struct $name_expr {
240            pub expr: $expr_type,
241        }
242
243        impl ExtensionExpr for $name_expr {
244            fn as_any(&self) -> &dyn Any {
245                self
246            }
247
248            fn name(&self) -> &str {
249                $extension_name
250            }
251
252            fn value_type(&self) -> ValueType {
253                self.expr.value_type()
254            }
255
256            fn children(&self) -> &[Expr] {
257                std::slice::from_ref(&self.expr)
258            }
259
260            fn with_new_children(&self, children: Vec<Expr>) -> Arc<dyn ExtensionExpr> {
261                let mut iter = children.into_iter();
262                match (iter.next(), iter.next()) {
263                    (Some(expr), None) => Arc::new($name_expr { expr }),
264                    _ => Arc::new(self.clone()),
265                }
266            }
267        }
268
269        #[allow(rustdoc::broken_intra_doc_links)]
270        #[derive(Debug, Clone)]
271        pub struct $name {
272            pub expr: Arc<$name_expr>,
273        }
274
275        impl $name {
276            pub fn new(expr: $expr_type) -> Self {
277                Self {
278                    expr: Arc::new($name_expr { expr }),
279                }
280            }
281        }
282    };
283}
284
285define_node_ast_extension!(Analyze, AnalyzeExpr, Expr, ANALYZE_NODE_NAME);
286define_node_ast_extension!(
287    AnalyzeVerbose,
288    AnalyzeVerboseExpr,
289    Expr,
290    ANALYZE_VERBOSE_NODE_NAME
291);
292define_node_ast_extension!(Explain, ExplainExpr, Expr, EXPLAIN_NODE_NAME);
293define_node_ast_extension!(
294    ExplainVerbose,
295    ExplainVerboseExpr,
296    Expr,
297    EXPLAIN_VERBOSE_NODE_NAME
298);
299#[derive(Debug, Clone)]
300pub struct AliasExpr {
301    pub expr: Expr,
302    pub alias: String,
303}
304impl ExtensionExpr for AliasExpr {
305    fn as_any(&self) -> &dyn Any {
306        self
307    }
308    fn name(&self) -> &str {
309        ALIAS_NODE_NAME
310    }
311    fn value_type(&self) -> ValueType {
312        self.expr.value_type()
313    }
314    fn children(&self) -> &[Expr] {
315        std::slice::from_ref(&self.expr)
316    }
317    fn with_new_children(&self, children: Vec<Expr>) -> Arc<dyn ExtensionExpr> {
318        let mut iter = children.into_iter();
319        match (iter.next(), iter.next()) {
320            (Some(expr), None) => Arc::new(Self {
321                expr,
322                alias: self.alias.clone(),
323            }),
324            _ => Arc::new(self.clone()),
325        }
326    }
327}
328#[derive(Debug, Clone)]
329pub struct Alias {
330    pub expr: Arc<AliasExpr>,
331}
332impl Alias {
333    pub fn new(expr: Expr, alias: String) -> Self {
334        Self {
335            expr: Arc::new(AliasExpr { expr, alias }),
336        }
337    }
338}
339
340#[cfg(test)]
341mod test {
342    use session::context::{QueryContext, QueryContextBuilder};
343
344    use super::*;
345
346    // Detailed logic tests are covered in the parser crate.
347    #[test]
348    fn parse_sql_simple() {
349        let sql = "select * from t1";
350        let stmt = QueryLanguageParser::parse_sql(sql, &QueryContext::arc()).unwrap();
351        let QueryStatement::Sql(sql_stmt) = stmt else {
352            panic!("Expected SQL statement, got {:?}", stmt);
353        };
354        assert_eq!("SELECT * FROM t1", sql_stmt.to_string());
355    }
356
357    #[test]
358    fn parse_sql_tql_uses_scheduled_time_extension() {
359        let ctx = Arc::new(
360            QueryContextBuilder::default()
361                .set_extension(
362                    crate::options::FLOW_SCHEDULED_TIME_MILLIS.to_string(),
363                    "1700000000000".to_string(),
364                )
365                .build(),
366        );
367        let query = "TQL EVAL (now() - '10 minutes'::interval, now(), '1m') http_requests_total";
368        let stmt = QueryLanguageParser::parse_sql(query, &ctx).unwrap();
369
370        match stmt {
371            QueryStatement::Sql(sql::statements::statement::Statement::Tql(
372                sql::statements::tql::Tql::Eval(eval),
373            )) => {
374                assert_eq!(eval.start, "1699999400");
375                assert_eq!(eval.end, "1700000000");
376                assert_eq!(eval.step, "1m");
377                assert_eq!(eval.query, "http_requests_total");
378            }
379            _ => panic!("Expected TQL eval statement, got {stmt:?}"),
380        }
381    }
382
383    #[test]
384    fn parse_promql_timestamp() {
385        let cases = vec![
386            (
387                "1435781451.781",
388                SystemTime::UNIX_EPOCH
389                    .checked_add(Duration::from_secs_f64(1435781451.781))
390                    .unwrap(),
391            ),
392            ("0.000", SystemTime::UNIX_EPOCH),
393            ("00", SystemTime::UNIX_EPOCH),
394            (
395                "2015-07-01T20:10:51.781Z",
396                SystemTime::UNIX_EPOCH
397                    .checked_add(Duration::from_secs_f64(1435781451.781))
398                    .unwrap(),
399            ),
400            ("1970-01-01T00:00:00.000Z", SystemTime::UNIX_EPOCH),
401        ];
402
403        for (input, expected) in cases {
404            let result = QueryLanguageParser::parse_promql_timestamp(input).unwrap();
405
406            let result = result
407                .duration_since(SystemTime::UNIX_EPOCH)
408                .unwrap()
409                .as_millis();
410            let expected = expected
411                .duration_since(SystemTime::UNIX_EPOCH)
412                .unwrap()
413                .as_millis();
414
415            // assert difference < 0.1 second
416            assert!(result.abs_diff(expected) < 100);
417        }
418
419        // i64::MAX + 1
420        let timestamp = "9223372036854775808.000";
421        let result = QueryLanguageParser::parse_promql_timestamp(timestamp);
422        assert_eq!(
423            result.unwrap_err().to_string(),
424            "Failed to add duration '9223372036854775808s' to SystemTime, overflowed"
425        );
426    }
427
428    #[test]
429    fn parse_promql_simple() {
430        let promql = PromQuery {
431            query: "http_request".to_string(),
432            start: "2022-02-13T17:14:00Z".to_string(),
433            end: "2023-02-13T17:14:00Z".to_string(),
434            step: "1d".to_string(),
435            lookback: "5m".to_string(),
436            alias: Some("my_query".to_string()),
437        };
438
439        #[cfg(not(windows))]
440        let expected = String::from(
441            "\
442            Promql(EvalStmt { \
443                expr: Extension(Extension { \
444                    expr: AliasExpr { \
445                        expr: VectorSelector(VectorSelector { \
446                            name: Some(\"http_request\"), \
447                            matchers: Matchers { matchers: [], or_matchers: [] }, \
448                            offset: None, at: None \
449                        }), \
450                        alias: \"my_query\" \
451                    } \
452                }), \
453                start: SystemTime { tv_sec: 1644772440, tv_nsec: 0 }, \
454                end: SystemTime { tv_sec: 1676308440, tv_nsec: 0 }, \
455                interval: 86400s, \
456                lookback_delta: 300s \
457            }, Some(\"my_query\"))",
458        );
459
460        // Windows has different debug output for SystemTime.
461        #[cfg(windows)]
462        let expected = String::from(
463            "\
464            Promql(EvalStmt { \
465                expr: Extension(Extension { \
466                    expr: AliasExpr { \
467                        expr: VectorSelector(VectorSelector { \
468                            name: Some(\"http_request\"), \
469                            matchers: Matchers { matchers: [], or_matchers: [] }, \
470                            offset: None, at: None \
471                        }), \
472                        alias: \"my_query\" \
473                    } \
474                }), \
475                start: SystemTime { intervals: 132892460400000000 }, \
476                end: SystemTime { intervals: 133207820400000000 }, \
477                interval: 86400s, \
478                lookback_delta: 300s \
479            }, Some(\"my_query\"))",
480        );
481
482        let result = QueryLanguageParser::parse_promql(&promql, &QueryContext::arc()).unwrap();
483        assert_eq!(format!("{result:?}"), expected);
484    }
485}
query/parser.rs

query/
parser.rs
