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};
17
18use sqlparser::ast::{Expr, ObjectName, Query, SetExpr, SqlOption, TableFactor, Value};
19
20use crate::error::{InvalidSqlSnafu, InvalidTableOptionValueSnafu, Result};
21
22/// Format an [ObjectName] without any quote of its idents.
23pub fn format_raw_object_name(name: &ObjectName) -> String {
24    struct Inner<'a> {
25        name: &'a ObjectName,
26    }
27
28    impl Display for Inner<'_> {
29        fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
30            let mut delim = "";
31            for ident in self.name.0.iter() {
32                write!(f, "{delim}")?;
33                delim = ".";
34                write!(f, "{}", ident.value)?;
35            }
36            Ok(())
37        }
38    }
39
40    format!("{}", Inner { name })
41}
42
43pub fn parse_option_string(option: SqlOption) -> Result<(String, String)> {
44    let SqlOption::KeyValue { key, value } = option else {
45        return InvalidSqlSnafu {
46            msg: "Expecting a key-value pair in the option",
47        }
48        .fail();
49    };
50    let v = match value {
51        Expr::Value(Value::SingleQuotedString(v)) | Expr::Value(Value::DoubleQuotedString(v)) => v,
52        Expr::Identifier(v) => v.value,
53        Expr::Value(Value::Number(v, _)) => v.to_string(),
54        value => return InvalidTableOptionValueSnafu { key, value }.fail(),
55    };
56    let k = key.value.to_lowercase();
57    Ok((k, v))
58}
59
60/// Walk through a [Query] and extract all the tables referenced in it.
61pub fn extract_tables_from_query(query: &Query) -> impl Iterator<Item = ObjectName> {
62    let mut names = HashSet::new();
63
64    extract_tables_from_set_expr(&query.body, &mut names);
65
66    names.into_iter()
67}
68
69/// Helper function for [extract_tables_from_query].
70///
71/// Handle [SetExpr].
72fn extract_tables_from_set_expr(set_expr: &SetExpr, names: &mut HashSet<ObjectName>) {
73    match set_expr {
74        SetExpr::Select(select) => {
75            for from in &select.from {
76                table_factor_to_object_name(&from.relation, names);
77                for join in &from.joins {
78                    table_factor_to_object_name(&join.relation, names);
79                }
80            }
81        }
82        SetExpr::Query(query) => {
83            extract_tables_from_set_expr(&query.body, names);
84        }
85        SetExpr::SetOperation { left, right, .. } => {
86            extract_tables_from_set_expr(left, names);
87            extract_tables_from_set_expr(right, names);
88        }
89        SetExpr::Values(_) | SetExpr::Insert(_) | SetExpr::Update(_) | SetExpr::Table(_) => {}
90    };
91}
92
93/// Helper function for [extract_tables_from_query].
94///
95/// Handle [TableFactor].
96fn table_factor_to_object_name(table_factor: &TableFactor, names: &mut HashSet<ObjectName>) {
97    if let TableFactor::Table { name, .. } = table_factor {
98        names.insert(name.to_owned());
99    }
100}
sql/util.rs

sql/
util.rs
