Skip to main content

databend_common_ast/parser/
parser.rs

1// Copyright 2021 Datafuse Labs
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 derive_visitor::DriveMut;
16use derive_visitor::VisitorMut;
17use nom::Parser;
18use pretty_assertions::assert_eq;
19
20use crate::ParseError;
21use crate::Range;
22use crate::Result;
23use crate::ast::DatabaseRef;
24use crate::ast::ExplainKind;
25use crate::ast::Expr;
26use crate::ast::Identifier;
27use crate::ast::Literal;
28use crate::ast::ProcedureIdentity;
29use crate::ast::SelectTarget;
30use crate::ast::Statement;
31use crate::ast::StatementWithFormat;
32use crate::ast::TableRef;
33use crate::parser::Backtrace;
34use crate::parser::common::IResult;
35use crate::parser::common::comma_separated_list0;
36use crate::parser::common::comma_separated_list1;
37use crate::parser::common::database_ref;
38use crate::parser::common::ident;
39use crate::parser::common::table_ref;
40use crate::parser::common::transform_span;
41use crate::parser::error::display_parser_error;
42use crate::parser::expr::expr;
43use crate::parser::expr::values;
44use crate::parser::input::Dialect;
45use crate::parser::input::Input;
46use crate::parser::input::ParseMode;
47use crate::parser::statement::insert_stmt;
48use crate::parser::statement::procedure_type_name;
49use crate::parser::statement::replace_stmt;
50use crate::parser::statement::statement;
51use crate::parser::token::Token;
52use crate::parser::token::TokenKind;
53use crate::parser::token::Tokenizer;
54
55pub fn tokenize_sql(sql: &str) -> Result<Vec<Token<'_>>> {
56    Tokenizer::new(sql).collect::<Result<Vec<_>>>()
57}
58
59/// Parse a SQL string into `Statement`s.
60#[fastrace::trace]
61pub fn parse_sql(tokens: &[Token], dialect: Dialect) -> Result<(Statement, Option<String>)> {
62    let stmt = run_parser(tokens, dialect, ParseMode::Default, false, statement)?;
63
64    #[cfg(debug_assertions)]
65    assert_reparse(tokens[0].source, stmt.clone());
66
67    Ok((stmt.stmt, stmt.format))
68}
69
70/// Parse udf function into Expr
71pub fn parse_expr(tokens: &[Token], dialect: Dialect) -> Result<Expr> {
72    run_parser(tokens, dialect, ParseMode::Default, false, expr)
73}
74
75/// Parse a table reference string like "table", "db.table", or "catalog.db.table".
76/// Correctly handles quoted identifiers like `"my.weird.table"`.
77pub fn parse_table_ref(sql: &str, dialect: Dialect) -> Result<TableRef> {
78    let tokens = tokenize_sql(sql)?;
79    run_parser(&tokens, dialect, ParseMode::Default, false, table_ref)
80}
81
82/// Parse a database reference string like "db" or "catalog.db".
83/// Correctly handles quoted identifiers.
84pub fn parse_database_ref(sql: &str, dialect: Dialect) -> Result<DatabaseRef> {
85    let tokens = tokenize_sql(sql)?;
86    run_parser(&tokens, dialect, ParseMode::Default, false, database_ref)
87}
88
89/// Parse a procedure reference string like "my_proc(INT, STRING)" or "my_proc()".
90/// Returns a `ProcedureIdentity` with the procedure name and argument types.
91pub fn parse_procedure_ref(sql: &str, dialect: Dialect) -> Result<ProcedureIdentity> {
92    let tokens = tokenize_sql(sql)?;
93    run_parser(&tokens, dialect, ParseMode::Default, false, |i| {
94        nom::combinator::map(
95            nom::sequence::pair(ident, procedure_type_name),
96            |(name, args_type): (Identifier, _)| ProcedureIdentity {
97                name: name.to_string(),
98                args_type,
99            },
100        )
101        .parse(i)
102    })
103}
104
105/// Parse a UDF name string (a single identifier).
106/// Rejects trailing tokens to avoid silently ignoring malformed input.
107pub fn parse_udf_ref(sql: &str, dialect: Dialect) -> Result<Identifier> {
108    let tokens = tokenize_sql(sql)?;
109    run_parser(&tokens, dialect, ParseMode::Default, false, ident)
110}
111
112pub fn parse_comma_separated_exprs(tokens: &[Token], dialect: Dialect) -> Result<Vec<Expr>> {
113    run_parser(tokens, dialect, ParseMode::Default, true, |i| {
114        comma_separated_list0(expr)(i)
115    })
116}
117
118pub fn parse_comma_separated_idents(tokens: &[Token], dialect: Dialect) -> Result<Vec<Identifier>> {
119    run_parser(tokens, dialect, ParseMode::Default, true, |i| {
120        comma_separated_list1(ident).parse(i)
121    })
122}
123
124pub fn parse_values(tokens: &[Token], dialect: Dialect) -> Result<Vec<Expr>> {
125    run_parser(tokens, dialect, ParseMode::Default, false, values)
126}
127
128pub fn parse_cluster_key_exprs(cluster_key: &str) -> Result<Vec<Expr>> {
129    // `cluster_key` is persisted in table metadata and may be created/rewritten under a
130    // different session dialect. Parse it with a dialect that accepts both identifier
131    // quote styles to keep ALTER behavior stable across sessions.
132    let tokens = tokenize_sql(cluster_key)?;
133    let mut ast_exprs = parse_comma_separated_exprs(&tokens, Dialect::default())?;
134    // unwrap tuple.
135    if ast_exprs.len() == 1
136        && let Expr::Tuple { exprs, .. } = &ast_exprs[0]
137    {
138        ast_exprs = exprs.clone();
139    }
140    Ok(ast_exprs)
141}
142
143pub fn parse_raw_insert_stmt(
144    tokens: &[Token],
145    dialect: Dialect,
146    in_streaming_load: bool,
147) -> Result<Statement> {
148    run_parser(
149        tokens,
150        dialect,
151        ParseMode::Default,
152        false,
153        insert_stmt(true, in_streaming_load),
154    )
155}
156
157pub fn parse_raw_replace_stmt(tokens: &[Token], dialect: Dialect) -> Result<Statement> {
158    run_parser(
159        tokens,
160        dialect,
161        ParseMode::Default,
162        false,
163        replace_stmt(true),
164    )
165}
166
167pub fn run_parser<O>(
168    tokens: &[Token],
169    dialect: Dialect,
170    mode: ParseMode,
171    allow_partial: bool,
172    mut parser: impl FnMut(Input) -> IResult<O>,
173) -> Result<O> {
174    let backtrace = Backtrace::new();
175    let input = Input {
176        tokens,
177        dialect,
178        mode,
179        backtrace: &backtrace,
180    };
181    match parser(input) {
182        Ok((rest, res)) => {
183            let is_complete = rest[0].kind == TokenKind::EOI;
184            if is_complete || allow_partial {
185                Ok(res)
186            } else {
187                Err(ParseError(
188                    transform_span(&rest[..1]),
189                    format!(
190                        "unable to parse rest of the sql, rest tokens:  {:?} ",
191                        rest.tokens
192                    ),
193                ))
194            }
195        }
196        Err(nom::Err::Error(err) | nom::Err::Failure(err)) => {
197            let source = tokens[0].source;
198            Err(ParseError(None, display_parser_error(err, source)))
199        }
200        Err(nom::Err::Incomplete(_)) => unreachable!(),
201    }
202}
203
204/// Check that the statement can be displayed and reparsed without loss
205#[allow(dead_code)]
206fn assert_reparse(sql: &str, stmt: StatementWithFormat) {
207    let stmt = reset_ast(stmt);
208
209    let new_sql = stmt.to_string();
210    let new_tokens = crate::parser::tokenize_sql(&new_sql).unwrap();
211    let new_stmt = run_parser(
212        &new_tokens,
213        Dialect::PostgreSQL,
214        ParseMode::Default,
215        false,
216        statement,
217    )
218    .map_err(|err| panic!("{} in {}", err.1, new_sql))
219    .unwrap();
220
221    let new_stmt = reset_ast(new_stmt);
222    assert_eq!(stmt, new_stmt, "\nleft:\n{}\nright:\n{}", sql, new_sql);
223}
224
225#[allow(dead_code)]
226fn reset_ast(mut stmt: StatementWithFormat) -> StatementWithFormat {
227    #[derive(VisitorMut)]
228    #[visitor(Range(enter), Literal(enter), ExplainKind(enter), SelectTarget(enter))]
229    struct ResetAST;
230
231    impl ResetAST {
232        fn enter_range(&mut self, range: &mut Range) {
233            range.start = 0;
234            range.end = 0;
235        }
236
237        fn enter_literal(&mut self, literal: &mut Literal) {
238            *literal = Literal::Null;
239        }
240
241        fn enter_explain_kind(&mut self, kind: &mut ExplainKind) {
242            match kind {
243                ExplainKind::Ast(_) => *kind = ExplainKind::Ast("".to_string()),
244                ExplainKind::Syntax(_) => *kind = ExplainKind::Syntax("".to_string()),
245                ExplainKind::Memo(_) => *kind = ExplainKind::Memo("".to_string()),
246                _ => (),
247            }
248        }
249
250        fn enter_select_target(&mut self, target: &mut SelectTarget) {
251            if let SelectTarget::StarColumns { column_filter, .. } = target {
252                *column_filter = None
253            }
254        }
255    }
256
257    stmt.drive_mut(&mut ResetAST);
258
259    stmt
260}