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 promql_parser::label::{METRIC_NAME, MatchOp};
20use promql_parser::parser::{
21    Expr as PromExpr, MatrixSelector as PromMatrixSelector, VectorSelector as PromVectorSelector,
22};
23use promql_parser::util::{ExprVisitor, walk_expr};
24use serde::Serialize;
25use snafu::ensure;
26use sqlparser::ast::{
27    Array, Expr, Ident, ObjectName, ObjectNamePart, SqlOption, Value, ValueWithSpan,
28    Visit as AstVisit, Visitor,
29};
30use sqlparser_derive::{Visit, VisitMut};
31
32use crate::ast::ObjectNamePartExt;
33use crate::error::{InvalidExprAsOptionValueSnafu, InvalidSqlSnafu, Result};
34use crate::parser::ParserContext;
35use crate::parsers::with_tql_parser::CteContent;
36use crate::statements::alter::AlterTableOperation;
37use crate::statements::comment::CommentObject;
38use crate::statements::copy::{Copy, CopyTable};
39use crate::statements::create::SqlOrTql;
40use crate::statements::query::Query;
41use crate::statements::statement::Statement;
42use crate::statements::tql::Tql;
43
44const SCHEMA_MATCHER: &str = "__schema__";
45const DATABASE_MATCHER: &str = "__database__";
46
47/// Format an [ObjectName] without any quote of its idents.
48pub fn format_raw_object_name(name: &ObjectName) -> String {
49    struct Inner<'a> {
50        name: &'a ObjectName,
51    }
52
53    impl Display for Inner<'_> {
54        fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
55            let mut delim = "";
56            for ident in self.name.0.iter() {
57                write!(f, "{delim}")?;
58                delim = ".";
59                write!(f, "{}", ident.to_string_unquoted())?;
60            }
61            Ok(())
62        }
63    }
64
65    format!("{}", Inner { name })
66}
67
68#[derive(Clone, Debug, PartialEq, Eq, Serialize, Visit, VisitMut)]
69pub struct OptionValue(Expr);
70
71impl OptionValue {
72    pub(crate) fn try_new(expr: Expr) -> Result<Self> {
73        ensure!(
74            matches!(
75                expr,
76                Expr::Value(_) | Expr::Identifier(_) | Expr::Array(_) | Expr::Struct { .. }
77            ),
78            InvalidExprAsOptionValueSnafu {
79                error: format!("{expr} not accepted")
80            }
81        );
82        Ok(Self(expr))
83    }
84
85    fn expr_as_string(expr: &Expr) -> Option<&str> {
86        match expr {
87            Expr::Value(ValueWithSpan { value, .. }) => match value {
88                Value::SingleQuotedString(s)
89                | Value::DoubleQuotedString(s)
90                | Value::TripleSingleQuotedString(s)
91                | Value::TripleDoubleQuotedString(s)
92                | Value::SingleQuotedByteStringLiteral(s)
93                | Value::DoubleQuotedByteStringLiteral(s)
94                | Value::TripleSingleQuotedByteStringLiteral(s)
95                | Value::TripleDoubleQuotedByteStringLiteral(s)
96                | Value::SingleQuotedRawStringLiteral(s)
97                | Value::DoubleQuotedRawStringLiteral(s)
98                | Value::TripleSingleQuotedRawStringLiteral(s)
99                | Value::TripleDoubleQuotedRawStringLiteral(s)
100                | Value::EscapedStringLiteral(s)
101                | Value::UnicodeStringLiteral(s)
102                | Value::NationalStringLiteral(s)
103                | Value::HexStringLiteral(s) => Some(s),
104                Value::DollarQuotedString(s) => Some(&s.value),
105                Value::Number(s, _) => Some(s),
106                Value::Boolean(b) => Some(if *b { "true" } else { "false" }),
107                _ => None,
108            },
109            Expr::Identifier(ident) => Some(&ident.value),
110            _ => None,
111        }
112    }
113
114    /// Convert the option value to a string.
115    ///
116    /// Notes: Not all values can be converted to a string, refer to [Self::expr_as_string] for more details.
117    pub fn as_string(&self) -> Option<&str> {
118        Self::expr_as_string(&self.0)
119    }
120
121    pub fn as_list(&self) -> Option<Vec<&str>> {
122        let expr = &self.0;
123        match expr {
124            Expr::Value(_) | Expr::Identifier(_) => self.as_string().map(|s| vec![s]),
125            Expr::Array(array) => array
126                .elem
127                .iter()
128                .map(Self::expr_as_string)
129                .collect::<Option<Vec<_>>>(),
130            _ => None,
131        }
132    }
133}
134
135impl From<String> for OptionValue {
136    fn from(value: String) -> Self {
137        Self(Expr::Identifier(Ident::new(value)))
138    }
139}
140
141impl From<&str> for OptionValue {
142    fn from(value: &str) -> Self {
143        Self(Expr::Identifier(Ident::new(value)))
144    }
145}
146
147impl From<Vec<&str>> for OptionValue {
148    fn from(value: Vec<&str>) -> Self {
149        Self(Expr::Array(Array {
150            elem: value
151                .into_iter()
152                .map(|x| Expr::Identifier(Ident::new(x)))
153                .collect(),
154            named: false,
155        }))
156    }
157}
158
159impl Display for OptionValue {
160    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
161        if let Some(s) = self.as_string() {
162            write!(f, "'{s}'")
163        } else if let Expr::Array(array) = &self.0
164            && array
165                .elem
166                .iter()
167                .all(|expr| Self::expr_as_string(expr).is_some())
168        {
169            f.write_str("[")?;
170            for (index, expr) in array.elem.iter().enumerate() {
171                if index > 0 {
172                    f.write_str(", ")?;
173                }
174                write!(f, "'{}'", Self::expr_as_string(expr).unwrap())?;
175            }
176            f.write_str("]")
177        } else {
178            write!(f, "'{}'", self.0)
179        }
180    }
181}
182
183pub fn parse_option_string(option: SqlOption) -> Result<(String, OptionValue)> {
184    let SqlOption::KeyValue { key, value } = option else {
185        return InvalidSqlSnafu {
186            msg: "Expecting a key-value pair in the option",
187        }
188        .fail();
189    };
190    let v = OptionValue::try_new(value)?;
191    let k = key.value.to_lowercase();
192    Ok((k, v))
193}
194
195/// Walk through a [Query] and extract all the tables referenced in it.
196pub fn extract_tables_from_query(query: &SqlOrTql) -> impl Iterator<Item = ObjectName> {
197    extract_tables_from_query_inner(query).0.into_iter()
198}
199
200/// Walk through a [Query] and extract its referenced tables, returning `None`
201/// if any table reference cannot be resolved statically.
202pub fn extract_tables_from_query_checked(
203    query: &SqlOrTql,
204) -> Option<impl Iterator<Item = ObjectName>> {
205    let (names, complete) = extract_tables_from_query_inner(query);
206    complete.then_some(names.into_iter())
207}
208
209/// Walks through a [`Statement`] and extracts all referenced tables, returning
210/// `None` if any table reference cannot be resolved statically.
211pub fn extract_tables_from_statement_checked(
212    stmt: &Statement,
213) -> Option<impl Iterator<Item = ObjectName>> {
214    let mut names = HashSet::new();
215    extract_tables_from_statement(stmt, &mut names).then_some(names.into_iter())
216}
217
218fn extract_tables_from_statement(stmt: &Statement, names: &mut HashSet<ObjectName>) -> bool {
219    match stmt {
220        Statement::Tql(tql) => extract_tables_from_tql(tql, names),
221        Statement::Insert(insert) => {
222            collect_relations(&insert.inner, names);
223            true
224        }
225        Statement::Delete(delete) => {
226            collect_relations(&delete.inner, names);
227            true
228        }
229        Statement::Query(query) => extract_tables_from_query_ast(query, names),
230        Statement::CreateTable(create) => {
231            names.insert(create.name.clone());
232            true
233        }
234        Statement::CreateExternalTable(create) => {
235            names.insert(create.name.clone());
236            true
237        }
238        Statement::AlterTable(alter) => {
239            let current = alter.table_name().clone();
240            names.insert(current.clone());
241            if let AlterTableOperation::RenameTable { new_table_name } = alter.alter_operation() {
242                let mut renamed = current;
243                if let Some(last) = renamed.0.last_mut() {
244                    *last = ObjectNamePart::Identifier(Ident::new(new_table_name));
245                }
246                names.insert(renamed);
247            }
248            true
249        }
250        Statement::DropTable(drop) => {
251            names.extend(drop.table_names().iter().cloned());
252            true
253        }
254        #[cfg(feature = "enterprise")]
255        Statement::UndropTable(undrop) => {
256            names.insert(undrop.table_name().clone());
257            true
258        }
259        Statement::TruncateTable(truncate) => {
260            names.insert(truncate.table_name().clone());
261            true
262        }
263        Statement::CreateTableLike(create) => {
264            names.extend([create.table_name.clone(), create.source_name.clone()]);
265            true
266        }
267        Statement::CreateView(create) => {
268            names.insert(create.name.clone());
269            extract_tables_from_statement(&create.query, names)
270        }
271        Statement::CreateFlow(create) => {
272            names.insert(create.sink_table_name.clone());
273            extract_tables_from_sql_or_tql(&create.query, names)
274        }
275        #[cfg(feature = "enterprise")]
276        Statement::CreateTrigger(create) => {
277            extract_tables_from_sql_or_tql(&create.trigger_on.query, names)
278        }
279        #[cfg(feature = "enterprise")]
280        Statement::AlterTrigger(alter) => match &alter.operation.trigger_on {
281            Some(trigger_on) => extract_tables_from_sql_or_tql(&trigger_on.query, names),
282            None => true,
283        },
284        Statement::DropView(drop) => {
285            names.insert(drop.view_name.clone());
286            true
287        }
288        Statement::Explain(explain) => extract_tables_from_statement(&explain.statement, names),
289        Statement::DeclareCursor(cursor) => extract_tables_from_query_ast(&cursor.query, names),
290        Statement::DescribeTable(describe) => {
291            names.insert(describe.name().clone());
292            true
293        }
294        Statement::ShowCreateTable(show) => {
295            names.insert(show.table_name.clone());
296            true
297        }
298        Statement::ShowCreateView(show) => {
299            names.insert(show.view_name.clone());
300            true
301        }
302        Statement::ShowColumns(show) => {
303            names.insert(metadata_table_name(&show.table, show.database.as_deref()));
304            true
305        }
306        Statement::ShowIndex(show) => {
307            names.insert(metadata_table_name(&show.table, show.database.as_deref()));
308            true
309        }
310        Statement::ShowRegion(show) => {
311            names.insert(metadata_table_name(&show.table, show.database.as_deref()));
312            true
313        }
314        Statement::Comment(comment) => {
315            if let CommentObject::Table(table) | CommentObject::Column { table, .. } =
316                &comment.object
317            {
318                names.insert(table.clone());
319            }
320            true
321        }
322        Statement::Copy(Copy::CopyTable(copy)) => {
323            let table = match copy {
324                CopyTable::To(args) | CopyTable::From(args) => &args.table_name,
325            };
326            names.insert(table.clone());
327            true
328        }
329        Statement::Copy(Copy::CopyQueryTo(copy)) => {
330            extract_tables_from_statement(&copy.query, names)
331        }
332        Statement::Copy(Copy::CopyDatabase(_)) => true,
333        Statement::DropDatabase(_)
334        | Statement::DropFlow(_)
335        | Statement::CreateDatabase(_)
336        | Statement::AlterDatabase(_)
337        | Statement::ShowDatabases(_)
338        | Statement::ShowTables(_)
339        | Statement::ShowTableStatus(_)
340        | Statement::ShowCharset(_)
341        | Statement::ShowCollation(_)
342        | Statement::ShowCreateDatabase(_)
343        | Statement::ShowCreateFlow(_)
344        | Statement::ShowFlows(_)
345        | Statement::ShowFlowStatus(_)
346        | Statement::ShowStatus(_)
347        | Statement::ShowSearchPath(_)
348        | Statement::ShowViews(_)
349        | Statement::SetVariables(_)
350        | Statement::ShowVariables(_)
351        | Statement::Use(_)
352        | Statement::Admin(_)
353        | Statement::FetchCursor(_)
354        | Statement::CloseCursor(_)
355        | Statement::Kill(_)
356        | Statement::ShowProcesslist(_) => true,
357        #[cfg(feature = "enterprise")]
358        Statement::DropTrigger(_)
359        | Statement::ShowCreateTrigger(_)
360        | Statement::ShowTriggers(_) => true,
361    }
362}
363
364fn extract_tables_from_query_ast(query: &Query, names: &mut HashSet<ObjectName>) -> bool {
365    extract_tables_from_sql_query(&query.inner, names);
366    extract_tables_from_hybrid_cte_query(query, names)
367}
368
369fn extract_tables_from_sql_or_tql(query: &SqlOrTql, names: &mut HashSet<ObjectName>) -> bool {
370    let (query_names, complete) = extract_tables_from_query_inner(query);
371    names.extend(query_names);
372    complete
373}
374
375fn collect_relations(ast: &impl AstVisit, names: &mut HashSet<ObjectName>) {
376    let _ = ast.visit(&mut RelationCollector::new(names));
377}
378
379fn metadata_table_name(table: &str, database: Option<&str>) -> ObjectName {
380    let mut parts = Vec::with_capacity(usize::from(database.is_some()) + 1);
381    parts.extend(database.map(Ident::new));
382    parts.push(Ident::new(table));
383    ObjectName::from(parts)
384}
385
386fn extract_tables_from_query_inner(query: &SqlOrTql) -> (HashSet<ObjectName>, bool) {
387    let mut names = HashSet::new();
388
389    let complete = match query {
390        SqlOrTql::Sql(query, _) => {
391            extract_tables_from_sql_query(&query.inner, &mut names);
392            extract_tables_from_hybrid_cte_query(query, &mut names)
393        }
394        SqlOrTql::Tql(tql, _) => extract_tables_from_tql(tql, &mut names),
395    };
396
397    (names, complete)
398}
399
400fn extract_tables_from_hybrid_cte_query(
401    query: &Query,
402    sql_names: &mut HashSet<ObjectName>,
403) -> bool {
404    let Some(hybrid_cte) = &query.hybrid_cte else {
405        return true;
406    };
407
408    let mut complete = true;
409    let cte_names: HashSet<String> = hybrid_cte
410        .cte_tables
411        .iter()
412        .map(|cte| ParserContext::canonicalize_identifier(cte.name.clone()).value)
413        .collect();
414    remove_cte_names(sql_names, &cte_names);
415
416    for cte in &hybrid_cte.cte_tables {
417        let mut cte_query_names = HashSet::new();
418        match &cte.content {
419            CteContent::Sql(cte_query) => {
420                extract_tables_from_sql_query(cte_query, &mut cte_query_names)
421            }
422            CteContent::Tql(tql) => complete &= extract_tables_from_tql(tql, &mut cte_query_names),
423        }
424        sql_names.extend(cte_query_names);
425    }
426
427    complete
428}
429
430fn remove_cte_names(names: &mut HashSet<ObjectName>, cte_names: &HashSet<String>) {
431    if cte_names.is_empty() {
432        return;
433    }
434
435    names.retain(|name| {
436        if name.0.len() != 1 {
437            return true;
438        }
439        let Some(ident) = name.0[0].as_ident() else {
440            return true;
441        };
442
443        let canonical = ParserContext::canonicalize_identifier(ident.clone()).value;
444        !cte_names.contains(&canonical)
445    });
446}
447
448fn extract_tables_from_tql(tql: &Tql, names: &mut HashSet<ObjectName>) -> bool {
449    let promql = match tql {
450        Tql::Eval(eval) => &eval.query,
451        Tql::Explain(explain) => &explain.query,
452        Tql::Analyze(analyze) => &analyze.query,
453    };
454
455    let Ok(expr) = promql_parser::parser::parse(promql) else {
456        return false;
457    };
458    extract_tables_from_prom_expr(&expr, names)
459}
460
461/// Extracts all tables referenced by a [`PromExpr`], returning `None`
462/// if any table reference cannot be resolved statically.
463pub fn extract_tables_from_prom_expr_checked(
464    expr: &PromExpr,
465) -> Option<impl Iterator<Item = ObjectName>> {
466    let mut names = HashSet::new();
467    extract_tables_from_prom_expr(expr, &mut names).then_some(names.into_iter())
468}
469
470fn extract_tables_from_prom_expr(expr: &PromExpr, names: &mut HashSet<ObjectName>) -> bool {
471    struct TableCollector<'a> {
472        names: &'a mut HashSet<ObjectName>,
473        complete: bool,
474    }
475
476    impl ExprVisitor for TableCollector<'_> {
477        type Error = ();
478
479        fn pre_visit(&mut self, expr: &PromExpr) -> std::result::Result<bool, Self::Error> {
480            self.complete &= match expr {
481                PromExpr::VectorSelector(selector) => {
482                    extract_metric_name_from_vector_selector(selector, self.names)
483                }
484                PromExpr::MatrixSelector(PromMatrixSelector { vs, .. }) => {
485                    extract_metric_name_from_vector_selector(vs, self.names)
486                }
487                PromExpr::Extension(_) => false,
488                _ => true,
489            };
490            Ok(true)
491        }
492    }
493
494    let mut collector = TableCollector {
495        names,
496        complete: true,
497    };
498    let _ = walk_expr(&mut collector, expr);
499    collector.complete
500}
501
502fn extract_metric_name_from_vector_selector(
503    selector: &PromVectorSelector,
504    names: &mut HashSet<ObjectName>,
505) -> bool {
506    if selector.name.is_none() && !selector.matchers.or_matchers.is_empty() {
507        return false;
508    }
509
510    let metric_name = selector.name.clone().or_else(|| {
511        let mut metric_name_matchers = selector.matchers.find_matchers(METRIC_NAME);
512        if metric_name_matchers.len() == 1 && metric_name_matchers[0].op == MatchOp::Equal {
513            metric_name_matchers.pop().map(|matcher| matcher.value)
514        } else {
515            None
516        }
517    });
518    let Some(metric_name) = metric_name else {
519        return false;
520    };
521
522    if selector.matchers.matchers.iter().any(|matcher| {
523        (matcher.name == SCHEMA_MATCHER || matcher.name == DATABASE_MATCHER)
524            && matcher.op != MatchOp::Equal
525    }) {
526        return false;
527    }
528
529    let schema_matcher = selector.matchers.matchers.iter().rev().find(|matcher| {
530        matcher.op == MatchOp::Equal
531            && (matcher.name == SCHEMA_MATCHER || matcher.name == DATABASE_MATCHER)
532    });
533
534    if let Some(schema) = schema_matcher {
535        names.insert(ObjectName(vec![
536            ObjectNamePart::Identifier(Ident::new(&schema.value)),
537            ObjectNamePart::Identifier(Ident::new(metric_name)),
538        ]));
539    } else {
540        names.insert(ObjectName(vec![ObjectNamePart::Identifier(Ident::new(
541            metric_name,
542        ))]));
543    }
544    true
545}
546
547/// translate the start location to the index in the sql string
548pub fn location_to_index(sql: &str, location: &sqlparser::tokenizer::Location) -> usize {
549    let mut index = 0;
550    for (lno, line) in sql.lines().enumerate() {
551        if lno + 1 == location.line as usize {
552            index += location.column as usize;
553            break;
554        } else {
555            index += line.len() + 1; // +1 for the newline
556        }
557    }
558    // -1 because the index is 0-based
559    // and the location is 1-based
560    index - 1
561}
562
563// Tracks CTE aliases as sqlparser's visitor enters their query bodies.
564struct QueryScope {
565    cte_names: Vec<String>,
566    next_cte: usize,
567    recursive: bool,
568    scope_start: usize,
569    add_to_parent_after: Option<String>,
570}
571
572struct RelationCollector<'a> {
573    names: &'a mut HashSet<ObjectName>,
574    ctes_in_scope: Vec<String>,
575    query_scopes: Vec<QueryScope>,
576    #[cfg(test)]
577    query_visits: usize,
578}
579
580impl<'a> RelationCollector<'a> {
581    fn new(names: &'a mut HashSet<ObjectName>) -> Self {
582        Self {
583            names,
584            ctes_in_scope: Vec::new(),
585            query_scopes: Vec::new(),
586            #[cfg(test)]
587            query_visits: 0,
588        }
589    }
590}
591
592impl Visitor for RelationCollector<'_> {
593    type Break = ();
594
595    fn pre_visit_query(&mut self, query: &sqlparser::ast::Query) -> ControlFlow<Self::Break> {
596        #[cfg(test)]
597        {
598            self.query_visits += 1;
599        }
600
601        // Query::visit enters WITH definitions in declaration order before its body.
602        let parent_cte = self.query_scopes.last_mut().and_then(|scope| {
603            let cte_name = scope.cte_names.get(scope.next_cte)?.clone();
604            scope.next_cte += 1;
605            Some((cte_name, scope.recursive))
606        });
607
608        let add_to_parent_after = match parent_cte {
609            Some((cte_name, true)) => {
610                self.ctes_in_scope.push(cte_name);
611                None
612            }
613            Some((cte_name, false)) => Some(cte_name),
614            None => None,
615        };
616
617        let scope_start = self.ctes_in_scope.len();
618        let (cte_names, recursive) = if let Some(with) = &query.with {
619            (
620                with.cte_tables
621                    .iter()
622                    .map(|cte| ParserContext::canonicalize_identifier(cte.alias.name.clone()).value)
623                    .collect(),
624                with.recursive,
625            )
626        } else {
627            (Vec::new(), false)
628        };
629
630        self.query_scopes.push(QueryScope {
631            cte_names,
632            next_cte: 0,
633            recursive,
634            scope_start,
635            add_to_parent_after,
636        });
637        ControlFlow::Continue(())
638    }
639
640    fn post_visit_query(&mut self, _query: &sqlparser::ast::Query) -> ControlFlow<Self::Break> {
641        let Some(scope) = self.query_scopes.pop() else {
642            return ControlFlow::Break(());
643        };
644        self.ctes_in_scope.truncate(scope.scope_start);
645        if let Some(cte_name) = scope.add_to_parent_after {
646            self.ctes_in_scope.push(cte_name);
647        }
648        ControlFlow::Continue(())
649    }
650
651    fn pre_visit_relation(&mut self, relation: &ObjectName) -> ControlFlow<Self::Break> {
652        let is_cte = matches!(
653            relation.0.as_slice(),
654            [part]
655                if part.as_ident().is_some_and(|ident| {
656                    self.ctes_in_scope.contains(
657                        &ParserContext::canonicalize_identifier(ident.clone()).value,
658                    )
659                })
660        );
661        if !is_cte {
662            self.names.insert(relation.clone());
663        }
664        ControlFlow::Continue(())
665    }
666}
667
668/// Helper function for [extract_tables_from_query].
669///
670/// Handle [sqlparser::ast::Query].
671fn extract_tables_from_sql_query(query: &sqlparser::ast::Query, names: &mut HashSet<ObjectName>) {
672    let _ = query.visit(&mut RelationCollector::new(names));
673}
674
675#[cfg(test)]
676mod tests {
677    use itertools::Itertools;
678    use sqlparser::tokenizer::Token;
679
680    use super::*;
681    use crate::dialect::GreptimeDbDialect;
682    use crate::parser::{ParseOptions, ParserContext};
683    use crate::statements::statement::Statement;
684
685    #[test]
686    fn test_location_to_index() {
687        let testcases = vec![
688            "SELECT * FROM t WHERE a = 1",
689            // start or end with newline
690            r"
691SELECT *
692FROM
693t
694WHERE a =
6951
696",
697            r"SELECT *
698FROM
699t
700WHERE a =
7011
702",
703            r"
704SELECT *
705FROM
706t
707WHERE a =
7081",
709        ];
710
711        for sql in testcases {
712            let mut parser = ParserContext::new(&GreptimeDbDialect {}, sql).unwrap();
713            loop {
714                let token = parser.parser.next_token();
715                if token == Token::EOF {
716                    break;
717                }
718                let span = token.span;
719                let subslice =
720                    &sql[location_to_index(sql, &span.start)..location_to_index(sql, &span.end)];
721                assert_eq!(token.to_string(), subslice);
722            }
723        }
724    }
725
726    #[test]
727    fn test_extract_tables_from_tql_query() {
728        let testcases = vec![
729            (
730                r#"
731CREATE FLOW calc_reqs SINK TO cnt_reqs AS
732TQL EVAL (now() - '15s'::interval, now(), '5s') count_values("status_code", http_requests);"#,
733                vec!["http_requests".to_string()],
734            ),
735            (
736                r#"
737CREATE FLOW calc_reqs SINK TO cnt_reqs AS
738TQL EVAL (now() - '15s'::interval, now(), '5s') count_values("status_code", {__name__="http_requests"});"#,
739                vec!["http_requests".to_string()],
740            ),
741        ];
742
743        for (sql, expected_tables) in testcases {
744            let mut stmts = ParserContext::create_with_dialect(
745                sql,
746                &GreptimeDbDialect {},
747                ParseOptions::default(),
748            )
749            .unwrap();
750            let Statement::CreateFlow(create_flow) = stmts.pop().unwrap() else {
751                unreachable!()
752            };
753
754            let mut tables = extract_tables_from_query(&create_flow.query)
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_tql_query_completeness() {
764        let testcases = [
765            ("cpu", true, vec!["cpu"]),
766            ("cpu + mem", true, vec!["cpu", "mem"]),
767            (r#"cpu{__database__="private"}"#, true, vec!["private.cpu"]),
768            ("1 + 2", true, vec![]),
769            (r#"{__name__="cpu"}"#, true, vec!["cpu"]),
770            (r#"cpu{job="api" or instance="host"}"#, true, vec!["cpu"]),
771            (r#"cpu + {job="api"}"#, false, vec!["cpu"]),
772            (r#"{__name__=~"cpu.*"}"#, false, vec![]),
773            (r#"{__name__=~"cpu.*" or job="api"}"#, false, vec![]),
774            (r#"cpu{__schema__=~"private.*"}"#, false, vec![]),
775        ];
776
777        for (promql, expected_complete, expected_tables) in testcases {
778            let expected_tables = expected_tables
779                .into_iter()
780                .map(str::to_string)
781                .collect_vec();
782            let sql = format!("TQL EVAL (0, 10, '5s') {promql}");
783            let mut stmts = ParserContext::create_with_dialect(
784                &sql,
785                &GreptimeDbDialect {},
786                ParseOptions::default(),
787            )
788            .unwrap();
789            let Statement::Tql(tql) = stmts.pop().unwrap() else {
790                unreachable!()
791            };
792
793            let query = SqlOrTql::Tql(tql, sql);
794            let (tables, complete) = extract_tables_from_query_inner(&query);
795            let mut tables = tables
796                .into_iter()
797                .map(|table| format_raw_object_name(&table))
798                .collect_vec();
799            tables.sort();
800            assert_eq!(expected_complete, complete, "{promql}");
801            assert_eq!(expected_tables, tables, "{promql}");
802            assert_eq!(
803                expected_complete,
804                extract_tables_from_query_checked(&query).is_some(),
805                "{promql}"
806            );
807
808            let expr = promql_parser::parser::parse(promql).unwrap();
809            let direct_tables = extract_tables_from_prom_expr_checked(&expr).map(|tables| {
810                let mut tables = tables
811                    .map(|table| format_raw_object_name(&table))
812                    .collect_vec();
813                tables.sort();
814                tables
815            });
816            assert_eq!(
817                expected_complete.then(|| expected_tables.clone()),
818                direct_tables
819            );
820        }
821    }
822
823    #[test]
824    fn test_extract_tables_from_statement() {
825        for (sql, expected) in [
826            ("SELECT * FROM physical_metric", vec!["physical_metric"]),
827            (
828                "INSERT INTO target SELECT * FROM physical_metric",
829                vec!["physical_metric", "target"],
830            ),
831            (
832                "INSERT INTO target WITH cte AS (SELECT * FROM physical_metric) SELECT * FROM cte",
833                vec!["physical_metric", "target"],
834            ),
835            (
836                "TQL EVAL (0, 10, '5s') physical_metric",
837                vec!["physical_metric"],
838            ),
839            (
840                "CREATE VIEW target AS SELECT * FROM physical_metric",
841                vec!["physical_metric", "target"],
842            ),
843            ("ALTER TABLE old RENAME new", vec!["new", "old"]),
844            #[cfg(feature = "enterprise")]
845            ("UNDROP TABLE restored", vec!["restored"]),
846            ("SHOW TABLES", vec![]),
847        ] {
848            let stmt = ParserContext::create_with_dialect(
849                sql,
850                &GreptimeDbDialect {},
851                ParseOptions::default(),
852            )
853            .unwrap()
854            .remove(0);
855            let mut tables = extract_tables_from_statement_checked(&stmt)
856                .unwrap()
857                .map(|table| format_raw_object_name(&table))
858                .collect_vec();
859            tables.sort();
860            assert_eq!(expected, tables, "{sql}");
861        }
862    }
863
864    #[test]
865    fn test_extract_tables_from_chained_tql_ctes() {
866        let sql = "WITH denied AS (TQL EVAL (0, 10, '5s') allowed), \
867                   leak AS (TQL EVAL (0, 10, '5s') denied) \
868                   SELECT * FROM leak";
869        let mut stmts =
870            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
871                .unwrap();
872        let Statement::Query(query) = stmts.pop().unwrap() else {
873            unreachable!()
874        };
875
876        let mut tables = extract_tables_from_query_checked(&SqlOrTql::Sql(*query, sql.to_string()))
877            .unwrap()
878            .map(|table| format_raw_object_name(&table))
879            .collect_vec();
880        tables.sort();
881        assert_eq!(vec!["allowed".to_string(), "denied".to_string()], tables);
882    }
883
884    #[test]
885    fn test_extract_tables_from_sql_query_with_derived_join() {
886        let sql = r#"
887CREATE FLOW flow_batch_join_subquery SINK TO flow_batch_join_sink
888EVAL INTERVAL '1m' AS
889SELECT a.symbol, b.mark_price
890FROM (
891    SELECT inst_id AS symbol, max(ts) AS mark_iv_ts
892    FROM flow_batch_join_opt_summary
893    GROUP BY inst_id
894) a
895LEFT JOIN (
896    SELECT symbol, max(mark_price) AS mark_price
897    FROM flow_batch_join_market_v5
898    WHERE "type" = 'OPTION_MARK'
899    GROUP BY symbol
900) b ON a.symbol = b.symbol;
901"#;
902        let mut stmts =
903            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
904                .unwrap();
905        let Statement::CreateFlow(create_flow) = stmts.pop().unwrap() else {
906            unreachable!()
907        };
908
909        let mut tables = extract_tables_from_query(&create_flow.query)
910            .map(|table| format_raw_object_name(&table))
911            .collect_vec();
912        tables.sort();
913        assert_eq!(
914            vec![
915                "flow_batch_join_market_v5".to_string(),
916                "flow_batch_join_opt_summary".to_string(),
917            ],
918            tables
919        );
920    }
921
922    #[test]
923    fn test_extract_tables_from_sql_query_with_expression_subqueries() {
924        let sql = r#"
925SELECT
926    (SELECT max(value) FROM scalar_source)
927FROM outer_source
928WHERE EXISTS (SELECT 1 FROM exists_source)
929ORDER BY (SELECT max(value) FROM order_source);
930"#;
931        let mut stmts =
932            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
933                .unwrap();
934        let Statement::Query(query) = stmts.pop().unwrap() else {
935            unreachable!()
936        };
937
938        let mut tables = extract_tables_from_query(&SqlOrTql::Sql(*query, sql.to_string()))
939            .map(|table| format_raw_object_name(&table))
940            .collect_vec();
941        tables.sort();
942        assert_eq!(
943            vec![
944                "exists_source".to_string(),
945                "order_source".to_string(),
946                "outer_source".to_string(),
947                "scalar_source".to_string(),
948            ],
949            tables
950        );
951    }
952
953    #[test]
954    fn test_extract_tables_from_sql_query_with_cte_scopes() {
955        let testcases = vec![
956            (
957                r#"
958WITH source AS (
959    SELECT * FROM source
960)
961SELECT * FROM source;
962"#,
963                vec!["source".to_string()],
964            ),
965            (
966                r#"
967WITH RECURSIVE source AS (
968    SELECT * FROM source
969)
970SELECT * FROM source;
971"#,
972                vec![],
973            ),
974            (
975                r#"
976WITH first_cte AS (
977    SELECT * FROM physical_source
978), second_cte AS (
979    SELECT * FROM first_cte
980)
981SELECT * FROM second_cte;
982"#,
983                vec!["physical_source".to_string()],
984            ),
985            (
986                r#"
987SELECT * FROM (
988    WITH nested_cte AS (SELECT * FROM nested_source)
989    SELECT * FROM nested_cte
990);
991"#,
992                vec!["nested_source".to_string()],
993            ),
994            (
995                r#"
996SELECT * FROM (
997    WITH nested_cte AS (SELECT * FROM nested_source)
998    SELECT (SELECT * FROM nested_cte)
999);
1000"#,
1001                vec!["nested_source".to_string()],
1002            ),
1003        ];
1004
1005        for (sql, expected_tables) in testcases {
1006            let mut stmts = ParserContext::create_with_dialect(
1007                sql,
1008                &GreptimeDbDialect {},
1009                ParseOptions::default(),
1010            )
1011            .unwrap();
1012            let Statement::Query(query) = stmts.pop().unwrap() else {
1013                unreachable!()
1014            };
1015
1016            let mut tables = HashSet::new();
1017            extract_tables_from_sql_query(&query.inner, &mut tables);
1018            let mut tables = tables
1019                .into_iter()
1020                .map(|table| format_raw_object_name(&table))
1021                .collect_vec();
1022            tables.sort();
1023            assert_eq!(expected_tables, tables);
1024        }
1025    }
1026
1027    #[test]
1028    fn test_extract_tables_from_deeply_nested_ctes_visits_each_query_once() {
1029        let mut sql = "SELECT * FROM physical_source".to_string();
1030        for depth in 0..20 {
1031            sql = format!("WITH cte_{depth} AS ({sql}) SELECT * FROM cte_{depth}");
1032        }
1033
1034        let mut stmts = ParserContext::create_with_dialect(
1035            &sql,
1036            &GreptimeDbDialect {},
1037            ParseOptions::default(),
1038        )
1039        .unwrap();
1040        let Statement::Query(query) = stmts.pop().unwrap() else {
1041            unreachable!()
1042        };
1043
1044        let mut tables = HashSet::new();
1045        let query_visits = {
1046            let mut collector = RelationCollector::new(&mut tables);
1047            let _ = query.inner.visit(&mut collector);
1048            collector.query_visits
1049        };
1050
1051        assert_eq!(21, query_visits);
1052        assert_eq!(
1053            vec!["physical_source"],
1054            tables
1055                .into_iter()
1056                .map(|table| format_raw_object_name(&table))
1057                .collect_vec()
1058        );
1059    }
1060
1061    #[test]
1062    fn test_extract_tables_from_tql_query_with_schema_matcher() {
1063        let sql = r#"
1064CREATE FLOW calc_reqs SINK TO cnt_reqs AS
1065TQL EVAL (now() - '15s'::interval, now(), '5s') count_values("status_code", http_requests{__schema__="greptime_private"});"#;
1066        let mut stmts =
1067            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1068                .unwrap();
1069        let Statement::CreateFlow(create_flow) = stmts.pop().unwrap() else {
1070            unreachable!()
1071        };
1072
1073        let mut tables = extract_tables_from_query(&create_flow.query)
1074            .map(|table| format_raw_object_name(&table))
1075            .collect_vec();
1076        tables.sort();
1077        assert_eq!(vec!["greptime_private.http_requests".to_string()], tables);
1078    }
1079}
sql/util.rs

sql/
util.rs
