feat: support multi-table insert

src/query/ast/src/ast/statements/insert_multi_table.rs

@@ -0,0 +1,130 @@
+// 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 std::fmt::Display;
+
+use derive_visitor::Drive;
+use derive_visitor::DriveMut;
+
+use crate::ast::write_comma_separated_list;
+use crate::ast::Expr;
+use crate::ast::Identifier;
+use crate::ast::Query;
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub struct IntoClause {
+    pub catalog: Option<Identifier>,
+    pub database: Option<Identifier>,
+    pub table: Identifier,
+    pub target_columns: Vec<Identifier>,
+    pub source_columns: Vec<Identifier>,
+}
+
+impl Display for IntoClause {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "INTO ")?;
+        if let Some(catalog) = &self.catalog {
+            write!(f, "{}.", catalog)?;
+        }
+        if let Some(database) = &self.database {
+            write!(f, "{}.", database)?;
+        }
+        write!(f, "{}", self.table)?;
+        if !self.target_columns.is_empty() {
+            write!(f, " (")?;
+            write_comma_separated_list(f, &self.target_columns)?;
+            write!(f, ")")?;
+        }
+        if !self.source_columns.is_empty() {
+            write!(f, " VALUES ")?;
+            write!(f, " (")?;
+            write_comma_separated_list(f, &self.source_columns)?;
+            write!(f, ")")?;
+        }
+        Ok(())
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub struct WhenClause {
+    pub condition: Expr,
+    pub into_clauses: Vec<IntoClause>,
+}
+
+impl Display for WhenClause {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "WHEN ")?;
+        self.condition.fmt(f)?;
+        write!(f, " THEN ")?;
+        for into_clause in &self.into_clauses {
+            write!(f, "{} ", into_clause)?;
+        }
+        Ok(())
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub struct ElseClause {
+    pub into_clauses: Vec<IntoClause>,
+}
+
+impl Display for ElseClause {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "ELSE ")?;
+        for into_clause in &self.into_clauses {
+            write!(f, "{} ", into_clause)?;
+        }
+        Ok(())
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub struct InsertMultiTableStmt {
+    #[drive(skip)]
+    pub overwrite: bool,
+    #[drive(skip)]
+    pub is_first: bool,
+    pub when_clauses: Vec<WhenClause>,
+    pub else_clause: Option<ElseClause>,
+    pub into_clauses: Vec<IntoClause>,
+    pub source: Query,
+}
+
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub enum InsertMultiTableKind {
+    First,
+    All,
+}
+
+impl Display for InsertMultiTableStmt {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "INSERT ")?;
+        if self.overwrite {
+            write!(f, "OVERWRITE ")?;
+        }
+        match &self.is_first {
+            true => write!(f, "FIRST ")?,
+            false => write!(f, "ALL ")?,
+        }
+        for when in &self.when_clauses {
+            write!(f, "{} ", when)?;
+        }
+        if let Some(else_clause) = &self.else_clause {
+            write!(f, "{} ", else_clause)?;
+        }
+        for into_clause in &self.into_clauses {
+            write!(f, "{} ", into_clause)?;
+        }
+        write!(f, "{}", self.source)
+    }
+}

src/query/ast/src/ast/statements/mod.rs

@@ -24,6 +24,7 @@ mod explain;
 mod hint;
 mod index;
 mod insert;
+mod insert_multi_table;
 mod kill;
 mod lock;
 mod merge_into;
@@ -60,6 +61,7 @@ pub use explain::*;
 pub use hint::*;
 pub use index::*;
 pub use insert::*;
+pub use insert_multi_table::*;
 pub use kill::*;
 pub use lock::*;
 pub use merge_into::*;

src/query/ast/src/ast/statements/statement.rs

@@ -104,6 +104,7 @@ pub enum Statement {
     },
 
     Insert(InsertStmt),
+    InsertMultiTable(InsertMultiTableStmt),
     Replace(ReplaceStmt),
     MergeInto(MergeIntoStmt),
     Delete(DeleteStmt),
@@ -415,6 +416,7 @@ impl Display for Statement {
             }
             Statement::Query(stmt) => write!(f, "{stmt}")?,
             Statement::Insert(stmt) => write!(f, "{stmt}")?,
+            Statement::InsertMultiTable(insert_multi_table) => write!(f, "{insert_multi_table}")?,
             Statement::Replace(stmt) => write!(f, "{stmt}")?,
             Statement::MergeInto(stmt) => write!(f, "{stmt}")?,
             Statement::Delete(stmt) => write!(f, "{stmt}")?,

src/query/ast/src/ast/visitors/walk.rs

@@ -567,5 +567,6 @@ pub fn walk_statement<'a, V: Visitor<'a>>(visitor: &mut V, statement: &'a Statem
         Statement::Begin => {}
         Statement::Commit => {}
         Statement::Abort => {}
+        Statement::InsertMultiTable(_) => {}
     }
 }

src/query/ast/src/ast/visitors/walk_mut.rs

@@ -565,5 +565,6 @@ pub fn walk_statement_mut<V: VisitorMut>(visitor: &mut V, statement: &mut Statem
         Statement::AlterNotification(stmt) => visitor.visit_alter_notification(stmt),
         Statement::DropNotification(stmt) => visitor.visit_drop_notification(stmt),
         Statement::DescribeNotification(stmt) => visitor.visit_describe_notification(stmt),
+        Statement::InsertMultiTable(_) => {}
     }
 }

src/query/ast/src/parser/statement.rs

@@ -2063,7 +2063,9 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
             | #show_password_policies: "`SHOW PASSWORD POLICIES [<show_options>]`"
         ),
         rule!(
-            #insert_stmt(false) : "`INSERT INTO [TABLE] <table> [(<column>, ...)] (FORMAT <format> | VALUES <values> | <query>)`"
+            #conditional_multi_table_insert() : "`INSERT [OVERWRITE] {FIRST|ALL} { WHEN <condition> THEN intoClause [ ... ] } [ ... ] [ ELSE intoClause ] <subquery>`"
+            | #unconditional_multi_table_insert() : "`INSERT [OVERWRITE] ALL intoClause [ ... ] <subquery>`"
+            | #insert_stmt(false) : "`INSERT INTO [TABLE] <table> [(<column>, ...)] (FORMAT <format> | VALUES <values> | <query>)`"
             | #replace_stmt(false) : "`REPLACE INTO [TABLE] <table> [(<column>, ...)] (FORMAT <format> | VALUES <values> | <query>)`"
             | #merge : "`MERGE INTO <target_table> USING <source> ON <join_expr> { matchedClause | notMatchedClause } [ ... ]`"
             | #delete : "`DELETE FROM <table> [WHERE ...]`"
@@ -2290,6 +2292,89 @@ pub fn insert_stmt(allow_raw: bool) -> impl FnMut(Input) -> IResult<Statement> {
     }
 }
 
+pub fn conditional_multi_table_insert() -> impl FnMut(Input) -> IResult<Statement> {
+    move |i| {
+        map(
+            rule! {
+                INSERT ~ OVERWRITE? ~ (FIRST | ALL) ~ (#when_clause)+ ~ (#else_clause)? ~ #query
+            },
+            |(_, overwrite, kind, when_clauses, opt_else, source)| {
+                Statement::InsertMultiTable(InsertMultiTableStmt {
+                    overwrite: overwrite.is_some(),
+                    is_first: matches!(kind.kind, FIRST),
+                    when_clauses,
+                    else_clause: opt_else,
+                    into_clauses: vec![],
+                    source,
+                })
+            },
+        )(i)
+    }
+}
+
+pub fn unconditional_multi_table_insert() -> impl FnMut(Input) -> IResult<Statement> {
+    move |i| {
+        map(
+            rule! {
+                INSERT ~ OVERWRITE? ~ ALL ~ (#into_clause)+ ~ #query
+            },
+            |(_, overwrite, _, into_clauses, source)| {
+                Statement::InsertMultiTable(InsertMultiTableStmt {
+                    overwrite: overwrite.is_some(),
+                    is_first: false,
+                    when_clauses: vec![],
+                    else_clause: None,
+                    into_clauses,
+                    source,
+                })
+            },
+        )(i)
+    }
+}
+
+fn when_clause(i: Input) -> IResult<WhenClause> {
+    map(
+        rule! {
+            WHEN ~ ^#expr ~ THEN ~ (#into_clause)+
+        },
+        |(_, expr, _, into_clauses)| WhenClause {
+            condition: expr,
+            into_clauses,
+        },
+    )(i)
+}
+
+fn into_clause(i: Input) -> IResult<IntoClause> {
+    map(
+        rule! {
+            INTO
+            ~ #dot_separated_idents_1_to_3
+            ~ ( "(" ~ #comma_separated_list1(ident) ~ ")" )?
+            ~ (VALUES ~ "(" ~ #comma_separated_list1(ident) ~ ")" )?
+        },
+        |(_, (catalog, database, table), opt_target_columns, opt_source_columns)| IntoClause {
+            catalog,
+            database,
+            table,
+            target_columns: opt_target_columns
+                .map(|(_, columns, _)| columns)
+                .unwrap_or_default(),
+            source_columns: opt_source_columns
+                .map(|(_, _, columns, _)| columns)
+                .unwrap_or_default(),
+        },
+    )(i)
+}
+
+fn else_clause(i: Input) -> IResult<ElseClause> {
+    map(
+        rule! {
+            ELSE ~ (#into_clause)+
+        },
+        |(_, into_clauses)| ElseClause { into_clauses },
+    )(i)
+}
+
 pub fn replace_stmt(allow_raw: bool) -> impl FnMut(Input) -> IResult<Statement> {
     move |i| {
         let insert_source_parser = if allow_raw {

src/query/ast/tests/it/display.rs

@@ -0,0 +1,46 @@
+// 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 databend_common_ast::parser::parse_sql;
+use databend_common_ast::parser::tokenize_sql;
+use databend_common_ast::parser::Dialect;
+
+fn test_stmt_display(sql: &str) {
+    let tokens = tokenize_sql(sql).unwrap();
+    let (stmt, _) = parse_sql(&tokens, Dialect::PostgreSQL).unwrap();
+    let sql1 = stmt.to_string();
+    let tokens1 = tokenize_sql(&sql1).unwrap();
+    let (stmt1, _) = parse_sql(&tokens1, Dialect::PostgreSQL).unwrap();
+    let sql2 = stmt1.to_string();
+    assert_eq!(sql1, sql2);
+}
+
+#[test]
+fn test_multi_table_insert_display() {
+    const SQL_FILE_PATH: &str = "tests/it/testsql/multi_table_insert.sql";
+    let sqls = std::fs::read_to_string(SQL_FILE_PATH).unwrap();
+    for sql in sqls.split(';').filter(|s| !s.is_empty()) {
+        test_stmt_display(sql);
+    }
+}
+
+#[test]
+fn test_multi_table_insert_parse_error() {
+    const SQL_FILE_PATH: &str = "tests/it/testsql/multi_table_insert_error.sql";
+    let sqls = std::fs::read_to_string(SQL_FILE_PATH).unwrap();
+    for sql in sqls.split(';').filter(|s| !s.is_empty()) {
+        let tokens = tokenize_sql(sql).unwrap();
+        assert!(parse_sql(&tokens, Dialect::PostgreSQL).is_err());
+    }
+}

src/query/ast/tests/it/main.rs

@@ -15,5 +15,6 @@
 #![allow(clippy::uninlined_format_args)]
 
 mod decimal;
+mod display;
 mod parser;
 mod token;

src/query/ast/tests/it/testsql/multi_table_insert.sql

@@ -0,0 +1,64 @@
+INSERT FIRST
+  WHEN n1 > 100 THEN
+    INTO t1
+  WHEN n1 > 10 THEN
+    INTO t1
+    INTO t2
+  ELSE
+    INTO t2
+SELECT n1 from src;
+
+INSERT OVERWRITE FIRST
+  WHEN n1 > 100 THEN
+    INTO t1
+  WHEN n1 > 10 THEN
+    INTO t1
+    INTO t2
+  ELSE
+    INTO t2
+SELECT n1 from src;
+
+INSERT OVERWRITE ALL
+  WHEN n1 > 100 THEN
+    INTO t1
+  WHEN n1 > 10 THEN
+    INTO t1
+    INTO t2
+  ELSE
+    INTO t2
+SELECT n1 from src;
+
+INSERT OVERWRITE ALL
+    INTO t1
+    INTO t2
+SELECT n1 from src;
+
+INSERT OVERWRITE FIRST
+  WHEN n1 > 100 THEN
+    INTO t1
+  WHEN n1 > 10 THEN
+    INTO t1
+    INTO t2
+SELECT n1 from src;
+
+INSERT OVERWRITE ALL
+  WHEN n1 > 100 THEN
+    INTO t1
+  WHEN n1 > 10 THEN
+    INTO t1
+    INTO t2
+SELECT n1 from src;
+
+INSERT ALL
+  INTO t1
+  INTO t1 (c1, c2, c3) VALUES (n2, n1, n3)
+  INTO t2 (c1, c2, c3)
+  INTO t2 VALUES (n3, n2, n1)
+SELECT n1, n2, n3 from src;
+
+INSERT OVERWRITE ALL
+  INTO t1
+  INTO t1 (c1, c2, c3) VALUES (n2, n1, n3)
+  INTO t2 (c1, c2, c3)
+  INTO t2 VALUES (n3, n2, n1)
+SELECT n1, n2, n3 from src;