databend-common-ast 0.2.5

// Copyright 2021 Datafuse Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use derive_visitor::DriveMut;
use derive_visitor::VisitorMut;
use nom::Parser;
use pretty_assertions::assert_eq;

use crate::ParseError;
use crate::Range;
use crate::Result;
use crate::ast::DatabaseRef;
use crate::ast::ExplainKind;
use crate::ast::Expr;
use crate::ast::Identifier;
use crate::ast::Literal;
use crate::ast::ProcedureIdentity;
use crate::ast::SelectTarget;
use crate::ast::Statement;
use crate::ast::StatementWithFormat;
use crate::ast::TableRef;
use crate::parser::Backtrace;
use crate::parser::common::IResult;
use crate::parser::common::comma_separated_list0;
use crate::parser::common::comma_separated_list1;
use crate::parser::common::database_ref;
use crate::parser::common::ident;
use crate::parser::common::table_ref;
use crate::parser::common::transform_span;
use crate::parser::error::display_parser_error;
use crate::parser::expr::expr;
use crate::parser::expr::values;
use crate::parser::input::Dialect;
use crate::parser::input::Input;
use crate::parser::input::ParseMode;
use crate::parser::statement::insert_stmt;
use crate::parser::statement::procedure_type_name;
use crate::parser::statement::replace_stmt;
use crate::parser::statement::statement;
use crate::parser::token::Token;
use crate::parser::token::TokenKind;
use crate::parser::token::Tokenizer;

pub fn tokenize_sql(sql: &str) -> Result<Vec<Token<'_>>> {
    Tokenizer::new(sql).collect::<Result<Vec<_>>>()
}

/// Parse a SQL string into `Statement`s.
#[fastrace::trace]
pub fn parse_sql(tokens: &[Token], dialect: Dialect) -> Result<(Statement, Option<String>)> {
    let stmt = run_parser(tokens, dialect, ParseMode::Default, false, statement)?;

    #[cfg(debug_assertions)]
    assert_reparse(tokens[0].source, stmt.clone());

    Ok((stmt.stmt, stmt.format))
}

/// Parse udf function into Expr
pub fn parse_expr(tokens: &[Token], dialect: Dialect) -> Result<Expr> {
    run_parser(tokens, dialect, ParseMode::Default, false, expr)
}

/// Parse a table reference string like "table", "db.table", or "catalog.db.table".
/// Correctly handles quoted identifiers like `"my.weird.table"`.
pub fn parse_table_ref(sql: &str, dialect: Dialect) -> Result<TableRef> {
    let tokens = tokenize_sql(sql)?;
    run_parser(&tokens, dialect, ParseMode::Default, false, table_ref)
}

/// Parse a database reference string like "db" or "catalog.db".
/// Correctly handles quoted identifiers.
pub fn parse_database_ref(sql: &str, dialect: Dialect) -> Result<DatabaseRef> {
    let tokens = tokenize_sql(sql)?;
    run_parser(&tokens, dialect, ParseMode::Default, false, database_ref)
}

/// Parse a procedure reference string like "my_proc(INT, STRING)" or "my_proc()".
/// Returns a `ProcedureIdentity` with the procedure name and argument types.
pub fn parse_procedure_ref(sql: &str, dialect: Dialect) -> Result<ProcedureIdentity> {
    let tokens = tokenize_sql(sql)?;
    run_parser(&tokens, dialect, ParseMode::Default, false, |i| {
        nom::combinator::map(
            nom::sequence::pair(ident, procedure_type_name),
            |(name, args_type): (Identifier, _)| ProcedureIdentity {
                name: name.to_string(),
                args_type,
            },
        )
        .parse(i)
    })
}

/// Parse a UDF name string (a single identifier).
/// Rejects trailing tokens to avoid silently ignoring malformed input.
pub fn parse_udf_ref(sql: &str, dialect: Dialect) -> Result<Identifier> {
    let tokens = tokenize_sql(sql)?;
    run_parser(&tokens, dialect, ParseMode::Default, false, ident)
}

pub fn parse_comma_separated_exprs(tokens: &[Token], dialect: Dialect) -> Result<Vec<Expr>> {
    run_parser(tokens, dialect, ParseMode::Default, true, |i| {
        comma_separated_list0(expr)(i)
    })
}

pub fn parse_comma_separated_idents(tokens: &[Token], dialect: Dialect) -> Result<Vec<Identifier>> {
    run_parser(tokens, dialect, ParseMode::Default, true, |i| {
        comma_separated_list1(ident).parse(i)
    })
}

pub fn parse_values(tokens: &[Token], dialect: Dialect) -> Result<Vec<Expr>> {
    run_parser(tokens, dialect, ParseMode::Default, false, values)
}

pub fn parse_cluster_key_exprs(cluster_key: &str) -> Result<Vec<Expr>> {
    // `cluster_key` is persisted in table metadata and may be created/rewritten under a
    // different session dialect. Parse it with a dialect that accepts both identifier
    // quote styles to keep ALTER behavior stable across sessions.
    let tokens = tokenize_sql(cluster_key)?;
    let mut ast_exprs = parse_comma_separated_exprs(&tokens, Dialect::default())?;
    // unwrap tuple.
    if ast_exprs.len() == 1
        && let Expr::Tuple { exprs, .. } = &ast_exprs[0]
    {
        ast_exprs = exprs.clone();
    }
    Ok(ast_exprs)
}

pub fn parse_raw_insert_stmt(
    tokens: &[Token],
    dialect: Dialect,
    in_streaming_load: bool,
) -> Result<Statement> {
    run_parser(
        tokens,
        dialect,
        ParseMode::Default,
        false,
        insert_stmt(true, in_streaming_load),
    )
}

pub fn parse_raw_replace_stmt(tokens: &[Token], dialect: Dialect) -> Result<Statement> {
    run_parser(
        tokens,
        dialect,
        ParseMode::Default,
        false,
        replace_stmt(true),
    )
}

pub fn run_parser<O>(
    tokens: &[Token],
    dialect: Dialect,
    mode: ParseMode,
    allow_partial: bool,
    mut parser: impl FnMut(Input) -> IResult<O>,
) -> Result<O> {
    let backtrace = Backtrace::new();
    let input = Input {
        tokens,
        dialect,
        mode,
        backtrace: &backtrace,
    };
    match parser(input) {
        Ok((rest, res)) => {
            let is_complete = rest[0].kind == TokenKind::EOI;
            if is_complete || allow_partial {
                Ok(res)
            } else {
                Err(ParseError(
                    transform_span(&rest[..1]),
                    format!(
                        "unable to parse rest of the sql, rest tokens:  {:?} ",
                        rest.tokens
                    ),
                ))
            }
        }
        Err(nom::Err::Error(err) | nom::Err::Failure(err)) => {
            let source = tokens[0].source;
            Err(ParseError(None, display_parser_error(err, source)))
        }
        Err(nom::Err::Incomplete(_)) => unreachable!(),
    }
}

/// Check that the statement can be displayed and reparsed without loss
#[allow(dead_code)]
fn assert_reparse(sql: &str, stmt: StatementWithFormat) {
    let stmt = reset_ast(stmt);

    let new_sql = stmt.to_string();
    let new_tokens = crate::parser::tokenize_sql(&new_sql).unwrap();
    let new_stmt = run_parser(
        &new_tokens,
        Dialect::PostgreSQL,
        ParseMode::Default,
        false,
        statement,
    )
    .map_err(|err| panic!("{} in {}", err.1, new_sql))
    .unwrap();

    let new_stmt = reset_ast(new_stmt);
    assert_eq!(stmt, new_stmt, "\nleft:\n{}\nright:\n{}", sql, new_sql);
}

#[allow(dead_code)]
fn reset_ast(mut stmt: StatementWithFormat) -> StatementWithFormat {
    #[derive(VisitorMut)]
    #[visitor(Range(enter), Literal(enter), ExplainKind(enter), SelectTarget(enter))]
    struct ResetAST;

    impl ResetAST {
        fn enter_range(&mut self, range: &mut Range) {
            range.start = 0;
            range.end = 0;
        }

        fn enter_literal(&mut self, literal: &mut Literal) {
            *literal = Literal::Null;
        }

        fn enter_explain_kind(&mut self, kind: &mut ExplainKind) {
            match kind {
                ExplainKind::Ast(_) => *kind = ExplainKind::Ast("".to_string()),
                ExplainKind::Syntax(_) => *kind = ExplainKind::Syntax("".to_string()),
                ExplainKind::Memo(_) => *kind = ExplainKind::Memo("".to_string()),
                _ => (),
            }
        }

        fn enter_select_target(&mut self, target: &mut SelectTarget) {
            if let SelectTarget::StarColumns { column_filter, .. } = target {
                *column_filter = None
            }
        }
    }

    stmt.drive_mut(&mut ResetAST);

    stmt
}