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

databend_common_ast/ast/statements/
merge_into.rs
