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

query/
parser.rs
