databend_common_ast/ast/statements/
merge_into.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 std::fmt::Display;
16use std::fmt::Formatter;
17
18use derive_visitor::Drive;
19use derive_visitor::DriveMut;
20
21use crate::ast::write_comma_separated_list;
22use crate::ast::write_dot_separated_list;
23use crate::ast::Expr;
24use crate::ast::Hint;
25use crate::ast::Identifier;
26use crate::ast::Query;
27use crate::ast::TableAlias;
28use crate::ast::TableReference;
29use crate::ast::WithOptions;
30
31#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
32pub struct MutationUpdateExpr {
33    pub table: Option<Identifier>,
34    pub name: Identifier,
35    pub expr: Expr,
36}
37
38impl Display for MutationUpdateExpr {
39    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
40        if self.table.is_some() {
41            write!(f, "{}.", self.table.clone().unwrap())?;
42        }
43
44        write!(f, "{} = {}", self.name, self.expr)
45    }
46}
47
48#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
49pub enum MatchOperation {
50    Update {
51        update_list: Vec<MutationUpdateExpr>,
52        is_star: bool,
53    },
54    Delete,
55}
56
57#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
58pub struct MatchedClause {
59    pub selection: Option<Expr>,
60    pub operation: MatchOperation,
61}
62
63#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
64pub struct InsertOperation {
65    pub columns: Option<Vec<Identifier>>,
66    pub values: Vec<Expr>,
67    pub is_star: bool,
68}
69
70#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
71pub struct UnmatchedClause {
72    pub selection: Option<Expr>,
73    pub insert_operation: InsertOperation,
74}
75
76#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
77pub enum MergeOption {
78    Match(MatchedClause),
79    Unmatch(UnmatchedClause),
80}
81
82#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
83pub struct MergeIntoStmt {
84    pub hints: Option<Hint>,
85    pub catalog: Option<Identifier>,
86    pub database: Option<Identifier>,
87    pub table_ident: Identifier,
88    pub source: MutationSource,
89    // target_alias is belong to target
90    pub target_alias: Option<TableAlias>,
91    pub join_expr: Expr,
92    pub merge_options: Vec<MergeOption>,
93}
94
95impl Display for MergeIntoStmt {
96    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
97        write!(f, "MERGE INTO ")?;
98        write_dot_separated_list(
99            f,
100            self.catalog
101                .iter()
102                .chain(&self.database)
103                .chain(Some(&self.table_ident)),
104        )?;
105        if let Some(alias) = &self.target_alias {
106            write!(f, " AS {}", alias.name)?;
107        }
108        write!(f, " USING {} ON {}", self.source, self.join_expr)?;
109
110        for clause in &self.merge_options {
111            match clause {
112                MergeOption::Match(match_clause) => {
113                    write!(f, " WHEN MATCHED ")?;
114                    if let Some(e) = &match_clause.selection {
115                        write!(f, "AND {} ", e)?;
116                    }
117                    write!(f, "THEN ")?;
118
119                    match &match_clause.operation {
120                        MatchOperation::Update {
121                            update_list,
122                            is_star,
123                        } => {
124                            if *is_star {
125                                write!(f, "UPDATE *")?;
126                            } else {
127                                write!(f, "UPDATE SET ")?;
128                                write_comma_separated_list(f, update_list)?;
129                            }
130                        }
131                        MatchOperation::Delete => {
132                            write!(f, "DELETE")?;
133                        }
134                    }
135                }
136                MergeOption::Unmatch(unmatch_clause) => {
137                    write!(f, " WHEN NOT MATCHED ")?;
138                    if let Some(e) = &unmatch_clause.selection {
139                        write!(f, "AND {} ", e)?;
140                    }
141                    write!(f, "THEN INSERT")?;
142
143                    if let Some(columns) = &unmatch_clause.insert_operation.columns {
144                        if !columns.is_empty() {
145                            write!(f, " (")?;
146                            write_comma_separated_list(f, columns)?;
147                            write!(f, ")")?;
148                        }
149                    }
150
151                    if unmatch_clause.insert_operation.is_star {
152                        write!(f, " *")?;
153                    } else {
154                        write!(f, " VALUES(")?;
155                        write_comma_separated_list(
156                            f,
157                            unmatch_clause.insert_operation.values.clone(),
158                        )?;
159                        write!(f, ")")?;
160                    }
161                }
162            }
163        }
164        Ok(())
165    }
166}
167
168#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
169pub enum MutationSource {
170    Select {
171        query: Box<Query>,
172        source_alias: TableAlias,
173    },
174    Table {
175        catalog: Option<Identifier>,
176        database: Option<Identifier>,
177        table: Identifier,
178        alias: Option<TableAlias>,
179        with_options: Option<WithOptions>,
180    },
181}
182
183impl MutationSource {
184    pub fn transform_table_reference(&self) -> TableReference {
185        match self {
186            Self::Select {
187                query,
188                source_alias,
189            } => TableReference::Subquery {
190                span: None,
191                lateral: false,
192                subquery: query.clone(),
193                alias: Some(source_alias.clone()),
194                pivot: None,
195                unpivot: None,
196            },
197            Self::Table {
198                catalog,
199                database,
200                table,
201                with_options,
202                alias,
203            } => TableReference::Table {
204                span: None,
205                catalog: catalog.clone(),
206                database: database.clone(),
207                table: table.clone(),
208                alias: alias.clone(),
209                temporal: None,
210                with_options: with_options.clone(),
211                pivot: None,
212                unpivot: None,
213                sample: None,
214            },
215        }
216    }
217}
218
219impl Display for MutationSource {
220    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
221        match self {
222            MutationSource::Select {
223                query,
224                source_alias,
225            } => write!(f, "({query}) AS {source_alias}"),
226
227            MutationSource::Table {
228                catalog,
229                database,
230                table,
231                with_options,
232                alias,
233            } => {
234                write_dot_separated_list(
235                    f,
236                    catalog.iter().chain(database.iter()).chain(Some(table)),
237                )?;
238                if let Some(with_options) = with_options {
239                    write!(f, " {with_options}")?;
240                }
241                if alias.is_some() {
242                    write!(f, " AS {}", alias.as_ref().unwrap())?;
243                }
244                Ok(())
245            }
246        }
247    }
248}
databend_common_ast/ast/statements/merge_into.rs

databend_common_ast/ast/statements/
merge_into.rs
