sql/
util.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::collections::HashSet;
16use std::fmt::{Display, Formatter};
17use std::ops::ControlFlow;
18
19use itertools::Itertools;
20use promql_parser::label::{METRIC_NAME, MatchOp};
21use promql_parser::parser::{
22    Expr as PromExpr, MatrixSelector as PromMatrixSelector, VectorSelector as PromVectorSelector,
23};
24use promql_parser::util::{ExprVisitor, walk_expr};
25use serde::Serialize;
26use snafu::ensure;
27use sqlparser::ast::{
28    Array, Expr, Ident, ObjectName, ObjectNamePart, SqlOption, Value, ValueWithSpan,
29    Visit as AstVisit, Visitor,
30};
31use sqlparser_derive::{Visit, VisitMut};
32
33use crate::ast::ObjectNamePartExt;
34use crate::error::{InvalidExprAsOptionValueSnafu, InvalidSqlSnafu, Result};
35use crate::parser::ParserContext;
36use crate::parsers::with_tql_parser::CteContent;
37use crate::statements::create::SqlOrTql;
38use crate::statements::query::Query;
39use crate::statements::tql::Tql;
40
41const SCHEMA_MATCHER: &str = "__schema__";
42const DATABASE_MATCHER: &str = "__database__";
43
44/// Format an [ObjectName] without any quote of its idents.
45pub fn format_raw_object_name(name: &ObjectName) -> String {
46    struct Inner<'a> {
47        name: &'a ObjectName,
48    }
49
50    impl Display for Inner<'_> {
51        fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
52            let mut delim = "";
53            for ident in self.name.0.iter() {
54                write!(f, "{delim}")?;
55                delim = ".";
56                write!(f, "{}", ident.to_string_unquoted())?;
57            }
58            Ok(())
59        }
60    }
61
62    format!("{}", Inner { name })
63}
64
65#[derive(Clone, Debug, PartialEq, Eq, Serialize, Visit, VisitMut)]
66pub struct OptionValue(Expr);
67
68impl OptionValue {
69    pub(crate) fn try_new(expr: Expr) -> Result<Self> {
70        ensure!(
71            matches!(
72                expr,
73                Expr::Value(_) | Expr::Identifier(_) | Expr::Array(_) | Expr::Struct { .. }
74            ),
75            InvalidExprAsOptionValueSnafu {
76                error: format!("{expr} not accepted")
77            }
78        );
79        Ok(Self(expr))
80    }
81
82    fn expr_as_string(expr: &Expr) -> Option<&str> {
83        match expr {
84            Expr::Value(ValueWithSpan { value, .. }) => match value {
85                Value::SingleQuotedString(s)
86                | Value::DoubleQuotedString(s)
87                | Value::TripleSingleQuotedString(s)
88                | Value::TripleDoubleQuotedString(s)
89                | Value::SingleQuotedByteStringLiteral(s)
90                | Value::DoubleQuotedByteStringLiteral(s)
91                | Value::TripleSingleQuotedByteStringLiteral(s)
92                | Value::TripleDoubleQuotedByteStringLiteral(s)
93                | Value::SingleQuotedRawStringLiteral(s)
94                | Value::DoubleQuotedRawStringLiteral(s)
95                | Value::TripleSingleQuotedRawStringLiteral(s)
96                | Value::TripleDoubleQuotedRawStringLiteral(s)
97                | Value::EscapedStringLiteral(s)
98                | Value::UnicodeStringLiteral(s)
99                | Value::NationalStringLiteral(s)
100                | Value::HexStringLiteral(s) => Some(s),
101                Value::DollarQuotedString(s) => Some(&s.value),
102                Value::Number(s, _) => Some(s),
103                Value::Boolean(b) => Some(if *b { "true" } else { "false" }),
104                _ => None,
105            },
106            Expr::Identifier(ident) => Some(&ident.value),
107            _ => None,
108        }
109    }
110
111    /// Convert the option value to a string.
112    ///
113    /// Notes: Not all values can be converted to a string, refer to [Self::expr_as_string] for more details.
114    pub fn as_string(&self) -> Option<&str> {
115        Self::expr_as_string(&self.0)
116    }
117
118    pub fn as_list(&self) -> Option<Vec<&str>> {
119        let expr = &self.0;
120        match expr {
121            Expr::Value(_) | Expr::Identifier(_) => self.as_string().map(|s| vec![s]),
122            Expr::Array(array) => array
123                .elem
124                .iter()
125                .map(Self::expr_as_string)
126                .collect::<Option<Vec<_>>>(),
127            _ => None,
128        }
129    }
130}
131
132impl From<String> for OptionValue {
133    fn from(value: String) -> Self {
134        Self(Expr::Identifier(Ident::new(value)))
135    }
136}
137
138impl From<&str> for OptionValue {
139    fn from(value: &str) -> Self {
140        Self(Expr::Identifier(Ident::new(value)))
141    }
142}
143
144impl From<Vec<&str>> for OptionValue {
145    fn from(value: Vec<&str>) -> Self {
146        Self(Expr::Array(Array {
147            elem: value
148                .into_iter()
149                .map(|x| Expr::Identifier(Ident::new(x)))
150                .collect(),
151            named: false,
152        }))
153    }
154}
155
156impl Display for OptionValue {
157    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
158        if let Some(s) = self.as_string() {
159            write!(f, "'{s}'")
160        } else if let Some(s) = self.as_list() {
161            write!(
162                f,
163                "[{}]",
164                s.into_iter().map(|x| format!("'{x}'")).join(", ")
165            )
166        } else {
167            write!(f, "'{}'", self.0)
168        }
169    }
170}
171
172pub fn parse_option_string(option: SqlOption) -> Result<(String, OptionValue)> {
173    let SqlOption::KeyValue { key, value } = option else {
174        return InvalidSqlSnafu {
175            msg: "Expecting a key-value pair in the option",
176        }
177        .fail();
178    };
179    let v = OptionValue::try_new(value)?;
180    let k = key.value.to_lowercase();
181    Ok((k, v))
182}
183
184/// Walk through a [Query] and extract all the tables referenced in it.
185pub fn extract_tables_from_query(query: &SqlOrTql) -> impl Iterator<Item = ObjectName> {
186    extract_tables_from_query_inner(query).0.into_iter()
187}
188
189/// Walk through a [Query] and extract its referenced tables, returning `None`
190/// if any table reference cannot be resolved statically.
191pub fn extract_tables_from_query_checked(
192    query: &SqlOrTql,
193) -> Option<impl Iterator<Item = ObjectName>> {
194    let (names, complete) = extract_tables_from_query_inner(query);
195    complete.then_some(names.into_iter())
196}
197
198fn extract_tables_from_query_inner(query: &SqlOrTql) -> (HashSet<ObjectName>, bool) {
199    let mut names = HashSet::new();
200
201    let complete = match query {
202        SqlOrTql::Sql(query, _) => {
203            extract_tables_from_sql_query(&query.inner, &mut names);
204            extract_tables_from_hybrid_cte_query(query, &mut names)
205        }
206        SqlOrTql::Tql(tql, _) => extract_tables_from_tql(tql, &mut names),
207    };
208
209    (names, complete)
210}
211
212fn extract_tables_from_hybrid_cte_query(
213    query: &Query,
214    sql_names: &mut HashSet<ObjectName>,
215) -> bool {
216    let Some(hybrid_cte) = &query.hybrid_cte else {
217        return true;
218    };
219
220    let mut complete = true;
221    let cte_names: HashSet<String> = hybrid_cte
222        .cte_tables
223        .iter()
224        .map(|cte| ParserContext::canonicalize_identifier(cte.name.clone()).value)
225        .collect();
226    remove_cte_names(sql_names, &cte_names);
227
228    for cte in &hybrid_cte.cte_tables {
229        let mut cte_query_names = HashSet::new();
230        match &cte.content {
231            CteContent::Sql(cte_query) => {
232                extract_tables_from_sql_query(cte_query, &mut cte_query_names)
233            }
234            CteContent::Tql(tql) => complete &= extract_tables_from_tql(tql, &mut cte_query_names),
235        }
236        sql_names.extend(cte_query_names);
237    }
238
239    complete
240}
241
242fn remove_cte_names(names: &mut HashSet<ObjectName>, cte_names: &HashSet<String>) {
243    if cte_names.is_empty() {
244        return;
245    }
246
247    names.retain(|name| {
248        if name.0.len() != 1 {
249            return true;
250        }
251        let Some(ident) = name.0[0].as_ident() else {
252            return true;
253        };
254
255        let canonical = ParserContext::canonicalize_identifier(ident.clone()).value;
256        !cte_names.contains(&canonical)
257    });
258}
259
260fn extract_tables_from_tql(tql: &Tql, names: &mut HashSet<ObjectName>) -> bool {
261    let promql = match tql {
262        Tql::Eval(eval) => &eval.query,
263        Tql::Explain(explain) => &explain.query,
264        Tql::Analyze(analyze) => &analyze.query,
265    };
266
267    let Ok(expr) = promql_parser::parser::parse(promql) else {
268        return false;
269    };
270    extract_tables_from_prom_expr(&expr, names)
271}
272
273fn extract_tables_from_prom_expr(expr: &PromExpr, names: &mut HashSet<ObjectName>) -> bool {
274    struct TableCollector<'a> {
275        names: &'a mut HashSet<ObjectName>,
276        complete: bool,
277    }
278
279    impl ExprVisitor for TableCollector<'_> {
280        type Error = ();
281
282        fn pre_visit(&mut self, expr: &PromExpr) -> std::result::Result<bool, Self::Error> {
283            self.complete &= match expr {
284                PromExpr::VectorSelector(selector) => {
285                    extract_metric_name_from_vector_selector(selector, self.names)
286                }
287                PromExpr::MatrixSelector(PromMatrixSelector { vs, .. }) => {
288                    extract_metric_name_from_vector_selector(vs, self.names)
289                }
290                PromExpr::Extension(_) => false,
291                _ => true,
292            };
293            Ok(true)
294        }
295    }
296
297    let mut collector = TableCollector {
298        names,
299        complete: true,
300    };
301    let _ = walk_expr(&mut collector, expr);
302    collector.complete
303}
304
305fn extract_metric_name_from_vector_selector(
306    selector: &PromVectorSelector,
307    names: &mut HashSet<ObjectName>,
308) -> bool {
309    let metric_name = selector.name.clone().or_else(|| {
310        let mut metric_name_matchers = selector.matchers.find_matchers(METRIC_NAME);
311        if metric_name_matchers.len() == 1 && metric_name_matchers[0].op == MatchOp::Equal {
312            metric_name_matchers.pop().map(|matcher| matcher.value)
313        } else {
314            None
315        }
316    });
317    let Some(metric_name) = metric_name else {
318        return false;
319    };
320
321    if selector.matchers.matchers.iter().any(|matcher| {
322        (matcher.name == SCHEMA_MATCHER || matcher.name == DATABASE_MATCHER)
323            && matcher.op != MatchOp::Equal
324    }) {
325        return false;
326    }
327
328    let schema_matcher = selector.matchers.matchers.iter().rev().find(|matcher| {
329        matcher.op == MatchOp::Equal
330            && (matcher.name == SCHEMA_MATCHER || matcher.name == DATABASE_MATCHER)
331    });
332
333    if let Some(schema) = schema_matcher {
334        names.insert(ObjectName(vec![
335            ObjectNamePart::Identifier(Ident::new(&schema.value)),
336            ObjectNamePart::Identifier(Ident::new(metric_name)),
337        ]));
338    } else {
339        names.insert(ObjectName(vec![ObjectNamePart::Identifier(Ident::new(
340            metric_name,
341        ))]));
342    }
343    true
344}
345
346/// translate the start location to the index in the sql string
347pub fn location_to_index(sql: &str, location: &sqlparser::tokenizer::Location) -> usize {
348    let mut index = 0;
349    for (lno, line) in sql.lines().enumerate() {
350        if lno + 1 == location.line as usize {
351            index += location.column as usize;
352            break;
353        } else {
354            index += line.len() + 1; // +1 for the newline
355        }
356    }
357    // -1 because the index is 0-based
358    // and the location is 1-based
359    index - 1
360}
361
362// Tracks CTE aliases as sqlparser's visitor enters their query bodies.
363struct QueryScope {
364    cte_names: Vec<String>,
365    next_cte: usize,
366    recursive: bool,
367    scope_start: usize,
368    add_to_parent_after: Option<String>,
369}
370
371struct RelationCollector<'a> {
372    names: &'a mut HashSet<ObjectName>,
373    ctes_in_scope: Vec<String>,
374    query_scopes: Vec<QueryScope>,
375    #[cfg(test)]
376    query_visits: usize,
377}
378
379impl<'a> RelationCollector<'a> {
380    fn new(names: &'a mut HashSet<ObjectName>) -> Self {
381        Self {
382            names,
383            ctes_in_scope: Vec::new(),
384            query_scopes: Vec::new(),
385            #[cfg(test)]
386            query_visits: 0,
387        }
388    }
389}
390
391impl Visitor for RelationCollector<'_> {
392    type Break = ();
393
394    fn pre_visit_query(&mut self, query: &sqlparser::ast::Query) -> ControlFlow<Self::Break> {
395        #[cfg(test)]
396        {
397            self.query_visits += 1;
398        }
399
400        // Query::visit enters WITH definitions in declaration order before its body.
401        let parent_cte = self.query_scopes.last_mut().and_then(|scope| {
402            let cte_name = scope.cte_names.get(scope.next_cte)?.clone();
403            scope.next_cte += 1;
404            Some((cte_name, scope.recursive))
405        });
406
407        let add_to_parent_after = match parent_cte {
408            Some((cte_name, true)) => {
409                self.ctes_in_scope.push(cte_name);
410                None
411            }
412            Some((cte_name, false)) => Some(cte_name),
413            None => None,
414        };
415
416        let scope_start = self.ctes_in_scope.len();
417        let (cte_names, recursive) = if let Some(with) = &query.with {
418            (
419                with.cte_tables
420                    .iter()
421                    .map(|cte| ParserContext::canonicalize_identifier(cte.alias.name.clone()).value)
422                    .collect(),
423                with.recursive,
424            )
425        } else {
426            (Vec::new(), false)
427        };
428
429        self.query_scopes.push(QueryScope {
430            cte_names,
431            next_cte: 0,
432            recursive,
433            scope_start,
434            add_to_parent_after,
435        });
436        ControlFlow::Continue(())
437    }
438
439    fn post_visit_query(&mut self, _query: &sqlparser::ast::Query) -> ControlFlow<Self::Break> {
440        let Some(scope) = self.query_scopes.pop() else {
441            return ControlFlow::Break(());
442        };
443        self.ctes_in_scope.truncate(scope.scope_start);
444        if let Some(cte_name) = scope.add_to_parent_after {
445            self.ctes_in_scope.push(cte_name);
446        }
447        ControlFlow::Continue(())
448    }
449
450    fn pre_visit_relation(&mut self, relation: &ObjectName) -> ControlFlow<Self::Break> {
451        let is_cte = matches!(
452            relation.0.as_slice(),
453            [part]
454                if part.as_ident().is_some_and(|ident| {
455                    self.ctes_in_scope.contains(
456                        &ParserContext::canonicalize_identifier(ident.clone()).value,
457                    )
458                })
459        );
460        if !is_cte {
461            self.names.insert(relation.clone());
462        }
463        ControlFlow::Continue(())
464    }
465}
466
467/// Helper function for [extract_tables_from_query].
468///
469/// Handle [sqlparser::ast::Query].
470fn extract_tables_from_sql_query(query: &sqlparser::ast::Query, names: &mut HashSet<ObjectName>) {
471    let _ = query.visit(&mut RelationCollector::new(names));
472}
473
474#[cfg(test)]
475mod tests {
476    use sqlparser::tokenizer::Token;
477
478    use super::*;
479    use crate::dialect::GreptimeDbDialect;
480    use crate::parser::{ParseOptions, ParserContext};
481    use crate::statements::statement::Statement;
482
483    #[test]
484    fn test_location_to_index() {
485        let testcases = vec![
486            "SELECT * FROM t WHERE a = 1",
487            // start or end with newline
488            r"
489SELECT *
490FROM
491t
492WHERE a =
4931
494",
495            r"SELECT *
496FROM
497t
498WHERE a =
4991
500",
501            r"
502SELECT *
503FROM
504t
505WHERE a =
5061",
507        ];
508
509        for sql in testcases {
510            let mut parser = ParserContext::new(&GreptimeDbDialect {}, sql).unwrap();
511            loop {
512                let token = parser.parser.next_token();
513                if token == Token::EOF {
514                    break;
515                }
516                let span = token.span;
517                let subslice =
518                    &sql[location_to_index(sql, &span.start)..location_to_index(sql, &span.end)];
519                assert_eq!(token.to_string(), subslice);
520            }
521        }
522    }
523
524    #[test]
525    fn test_extract_tables_from_tql_query() {
526        let testcases = vec![
527            (
528                r#"
529CREATE FLOW calc_reqs SINK TO cnt_reqs AS
530TQL EVAL (now() - '15s'::interval, now(), '5s') count_values("status_code", http_requests);"#,
531                vec!["http_requests".to_string()],
532            ),
533            (
534                r#"
535CREATE FLOW calc_reqs SINK TO cnt_reqs AS
536TQL EVAL (now() - '15s'::interval, now(), '5s') count_values("status_code", {__name__="http_requests"});"#,
537                vec!["http_requests".to_string()],
538            ),
539        ];
540
541        for (sql, expected_tables) in testcases {
542            let mut stmts = ParserContext::create_with_dialect(
543                sql,
544                &GreptimeDbDialect {},
545                ParseOptions::default(),
546            )
547            .unwrap();
548            let Statement::CreateFlow(create_flow) = stmts.pop().unwrap() else {
549                unreachable!()
550            };
551
552            let mut tables = extract_tables_from_query(&create_flow.query)
553                .map(|table| format_raw_object_name(&table))
554                .collect_vec();
555            tables.sort();
556            assert_eq!(expected_tables, tables);
557        }
558    }
559
560    #[test]
561    fn test_extract_tables_from_tql_query_completeness() {
562        let testcases = [
563            ("cpu", true, vec!["cpu"]),
564            (r#"{__name__="cpu"}"#, true, vec!["cpu"]),
565            (r#"cpu + {job="api"}"#, false, vec!["cpu"]),
566            (r#"{__name__=~"cpu.*"}"#, false, vec![]),
567            (r#"cpu{__schema__=~"private.*"}"#, false, vec![]),
568        ];
569
570        for (promql, expected_complete, expected_tables) in testcases {
571            let sql = format!("TQL EVAL (0, 10, '5s') {promql}");
572            let mut stmts = ParserContext::create_with_dialect(
573                &sql,
574                &GreptimeDbDialect {},
575                ParseOptions::default(),
576            )
577            .unwrap();
578            let Statement::Tql(tql) = stmts.pop().unwrap() else {
579                unreachable!()
580            };
581
582            let query = SqlOrTql::Tql(tql, sql);
583            let (tables, complete) = extract_tables_from_query_inner(&query);
584            let mut tables = tables
585                .into_iter()
586                .map(|table| format_raw_object_name(&table))
587                .collect_vec();
588            tables.sort();
589            assert_eq!(expected_complete, complete, "{promql}");
590            assert_eq!(expected_tables, tables, "{promql}");
591            assert_eq!(
592                expected_complete,
593                extract_tables_from_query_checked(&query).is_some(),
594                "{promql}"
595            );
596        }
597    }
598
599    #[test]
600    fn test_extract_tables_from_chained_tql_ctes() {
601        let sql = "WITH denied AS (TQL EVAL (0, 10, '5s') allowed), \
602                   leak AS (TQL EVAL (0, 10, '5s') denied) \
603                   SELECT * FROM leak";
604        let mut stmts =
605            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
606                .unwrap();
607        let Statement::Query(query) = stmts.pop().unwrap() else {
608            unreachable!()
609        };
610
611        let mut tables = extract_tables_from_query_checked(&SqlOrTql::Sql(*query, sql.to_string()))
612            .unwrap()
613            .map(|table| format_raw_object_name(&table))
614            .collect_vec();
615        tables.sort();
616        assert_eq!(vec!["allowed".to_string(), "denied".to_string()], tables);
617    }
618
619    #[test]
620    fn test_extract_tables_from_sql_query_with_derived_join() {
621        let sql = r#"
622CREATE FLOW flow_batch_join_subquery SINK TO flow_batch_join_sink
623EVAL INTERVAL '1m' AS
624SELECT a.symbol, b.mark_price
625FROM (
626    SELECT inst_id AS symbol, max(ts) AS mark_iv_ts
627    FROM flow_batch_join_opt_summary
628    GROUP BY inst_id
629) a
630LEFT JOIN (
631    SELECT symbol, max(mark_price) AS mark_price
632    FROM flow_batch_join_market_v5
633    WHERE "type" = 'OPTION_MARK'
634    GROUP BY symbol
635) b ON a.symbol = b.symbol;
636"#;
637        let mut stmts =
638            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
639                .unwrap();
640        let Statement::CreateFlow(create_flow) = stmts.pop().unwrap() else {
641            unreachable!()
642        };
643
644        let mut tables = extract_tables_from_query(&create_flow.query)
645            .map(|table| format_raw_object_name(&table))
646            .collect_vec();
647        tables.sort();
648        assert_eq!(
649            vec![
650                "flow_batch_join_market_v5".to_string(),
651                "flow_batch_join_opt_summary".to_string(),
652            ],
653            tables
654        );
655    }
656
657    #[test]
658    fn test_extract_tables_from_sql_query_with_expression_subqueries() {
659        let sql = r#"
660SELECT
661    (SELECT max(value) FROM scalar_source)
662FROM outer_source
663WHERE EXISTS (SELECT 1 FROM exists_source)
664ORDER BY (SELECT max(value) FROM order_source);
665"#;
666        let mut stmts =
667            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
668                .unwrap();
669        let Statement::Query(query) = stmts.pop().unwrap() else {
670            unreachable!()
671        };
672
673        let mut tables = extract_tables_from_query(&SqlOrTql::Sql(*query, sql.to_string()))
674            .map(|table| format_raw_object_name(&table))
675            .collect_vec();
676        tables.sort();
677        assert_eq!(
678            vec![
679                "exists_source".to_string(),
680                "order_source".to_string(),
681                "outer_source".to_string(),
682                "scalar_source".to_string(),
683            ],
684            tables
685        );
686    }
687
688    #[test]
689    fn test_extract_tables_from_sql_query_with_cte_scopes() {
690        let testcases = vec![
691            (
692                r#"
693WITH source AS (
694    SELECT * FROM source
695)
696SELECT * FROM source;
697"#,
698                vec!["source".to_string()],
699            ),
700            (
701                r#"
702WITH RECURSIVE source AS (
703    SELECT * FROM source
704)
705SELECT * FROM source;
706"#,
707                vec![],
708            ),
709            (
710                r#"
711WITH first_cte AS (
712    SELECT * FROM physical_source
713), second_cte AS (
714    SELECT * FROM first_cte
715)
716SELECT * FROM second_cte;
717"#,
718                vec!["physical_source".to_string()],
719            ),
720            (
721                r#"
722SELECT * FROM (
723    WITH nested_cte AS (SELECT * FROM nested_source)
724    SELECT * FROM nested_cte
725);
726"#,
727                vec!["nested_source".to_string()],
728            ),
729            (
730                r#"
731SELECT * FROM (
732    WITH nested_cte AS (SELECT * FROM nested_source)
733    SELECT (SELECT * FROM nested_cte)
734);
735"#,
736                vec!["nested_source".to_string()],
737            ),
738        ];
739
740        for (sql, expected_tables) in testcases {
741            let mut stmts = ParserContext::create_with_dialect(
742                sql,
743                &GreptimeDbDialect {},
744                ParseOptions::default(),
745            )
746            .unwrap();
747            let Statement::Query(query) = stmts.pop().unwrap() else {
748                unreachable!()
749            };
750
751            let mut tables = HashSet::new();
752            extract_tables_from_sql_query(&query.inner, &mut tables);
753            let mut tables = tables
754                .into_iter()
755                .map(|table| format_raw_object_name(&table))
756                .collect_vec();
757            tables.sort();
758            assert_eq!(expected_tables, tables);
759        }
760    }
761
762    #[test]
763    fn test_extract_tables_from_deeply_nested_ctes_visits_each_query_once() {
764        let mut sql = "SELECT * FROM physical_source".to_string();
765        for depth in 0..20 {
766            sql = format!("WITH cte_{depth} AS ({sql}) SELECT * FROM cte_{depth}");
767        }
768
769        let mut stmts = ParserContext::create_with_dialect(
770            &sql,
771            &GreptimeDbDialect {},
772            ParseOptions::default(),
773        )
774        .unwrap();
775        let Statement::Query(query) = stmts.pop().unwrap() else {
776            unreachable!()
777        };
778
779        let mut tables = HashSet::new();
780        let query_visits = {
781            let mut collector = RelationCollector::new(&mut tables);
782            let _ = query.inner.visit(&mut collector);
783            collector.query_visits
784        };
785
786        assert_eq!(21, query_visits);
787        assert_eq!(
788            vec!["physical_source"],
789            tables
790                .into_iter()
791                .map(|table| format_raw_object_name(&table))
792                .collect_vec()
793        );
794    }
795
796    #[test]
797    fn test_extract_tables_from_tql_query_with_schema_matcher() {
798        let sql = r#"
799CREATE FLOW calc_reqs SINK TO cnt_reqs AS
800TQL EVAL (now() - '15s'::interval, now(), '5s') count_values("status_code", http_requests{__schema__="greptime_private"});"#;
801        let mut stmts =
802            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
803                .unwrap();
804        let Statement::CreateFlow(create_flow) = stmts.pop().unwrap() else {
805            unreachable!()
806        };
807
808        let mut tables = extract_tables_from_query(&create_flow.query)
809            .map(|table| format_raw_object_name(&table))
810            .collect_vec();
811        tables.sort();
812        assert_eq!(vec!["greptime_private.http_requests".to_string()], tables);
813    }
814}
sql/util.rs

sql/
util.rs
