sqlparser/parser/
merge.rs

1// Licensed under the Apache License, Version 2.0 (the "License");
2// you may not use this file except in compliance with the License.
3// You may obtain a copy of the License at
4//
5// http://www.apache.org/licenses/LICENSE-2.0
6//
7// Unless required by applicable law or agreed to in writing, software
8// distributed under the License is distributed on an "AS IS" BASIS,
9// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10// See the License for the specific language governing permissions and
11// limitations under the License.
12
13//! SQL Parser for a `MERGE` statement
14
15#[cfg(not(feature = "std"))]
16use alloc::{boxed::Box, format, string::ToString, vec, vec::Vec};
17
18use crate::{
19    ast::{
20        Merge, MergeAction, MergeClause, MergeClauseKind, MergeInsertExpr, MergeInsertKind,
21        MergeUpdateExpr, ObjectName, OutputClause, SetExpr, Statement,
22    },
23    dialect::{BigQueryDialect, GenericDialect, MySqlDialect},
24    keywords::Keyword,
25    parser::IsOptional,
26    tokenizer::TokenWithSpan,
27};
28
29use super::{Parser, ParserError};
30
31impl Parser<'_> {
32    /// Parse a `MERGE` statement, returning a `Box`ed SetExpr
33    ///
34    /// This is used to reduce the size of the stack frames in debug builds
35    pub(super) fn parse_merge_setexpr_boxed(
36        &mut self,
37        merge_token: TokenWithSpan,
38    ) -> Result<Box<SetExpr>, ParserError> {
39        Ok(Box::new(SetExpr::Merge(self.parse_merge(merge_token)?)))
40    }
41
42    pub fn parse_merge(&mut self, merge_token: TokenWithSpan) -> Result<Statement, ParserError> {
43        let into = self.parse_keyword(Keyword::INTO);
44
45        let table = self.parse_table_factor()?;
46
47        self.expect_keyword_is(Keyword::USING)?;
48        let source = self.parse_table_factor()?;
49        self.expect_keyword_is(Keyword::ON)?;
50        let on = self.parse_expr()?;
51        let clauses = self.parse_merge_clauses()?;
52        let output = match self.parse_one_of_keywords(&[Keyword::OUTPUT, Keyword::RETURNING]) {
53            Some(keyword) => Some(self.parse_output(keyword, self.get_current_token().clone())?),
54            None => None,
55        };
56
57        Ok(Statement::Merge(Merge {
58            merge_token: merge_token.into(),
59            into,
60            table,
61            source,
62            on: Box::new(on),
63            clauses,
64            output,
65        }))
66    }
67
68    fn parse_merge_clauses(&mut self) -> Result<Vec<MergeClause>, ParserError> {
69        let mut clauses = vec![];
70        loop {
71            if !(self.parse_keyword(Keyword::WHEN)) {
72                break;
73            }
74            let when_token = self.get_current_token().clone();
75
76            let mut clause_kind = MergeClauseKind::Matched;
77            if self.parse_keyword(Keyword::NOT) {
78                clause_kind = MergeClauseKind::NotMatched;
79            }
80            self.expect_keyword_is(Keyword::MATCHED)?;
81
82            if matches!(clause_kind, MergeClauseKind::NotMatched)
83                && self.parse_keywords(&[Keyword::BY, Keyword::SOURCE])
84            {
85                clause_kind = MergeClauseKind::NotMatchedBySource;
86            } else if matches!(clause_kind, MergeClauseKind::NotMatched)
87                && self.parse_keywords(&[Keyword::BY, Keyword::TARGET])
88            {
89                clause_kind = MergeClauseKind::NotMatchedByTarget;
90            }
91
92            let predicate = if self.parse_keyword(Keyword::AND) {
93                Some(self.parse_expr()?)
94            } else {
95                None
96            };
97
98            self.expect_keyword_is(Keyword::THEN)?;
99
100            let merge_clause = match self.parse_one_of_keywords(&[
101                Keyword::UPDATE,
102                Keyword::INSERT,
103                Keyword::DELETE,
104            ]) {
105                Some(Keyword::UPDATE) => {
106                    if matches!(
107                        clause_kind,
108                        MergeClauseKind::NotMatched | MergeClauseKind::NotMatchedByTarget
109                    ) {
110                        return parser_err!(
111                            format_args!("UPDATE is not allowed in a {clause_kind} merge clause"),
112                            self.get_current_token().span.start
113                        );
114                    }
115
116                    let update_token = self.get_current_token().clone();
117                    self.expect_keyword_is(Keyword::SET)?;
118                    let assignments = self.parse_comma_separated(Parser::parse_assignment)?;
119                    let update_predicate = if self.parse_keyword(Keyword::WHERE) {
120                        Some(self.parse_expr()?)
121                    } else {
122                        None
123                    };
124                    let delete_predicate = if self.parse_keyword(Keyword::DELETE) {
125                        let _ = self.expect_keyword(Keyword::WHERE)?;
126                        Some(self.parse_expr()?)
127                    } else {
128                        None
129                    };
130                    MergeAction::Update(MergeUpdateExpr {
131                        update_token: update_token.into(),
132                        assignments,
133                        update_predicate,
134                        delete_predicate,
135                    })
136                }
137                Some(Keyword::DELETE) => {
138                    if matches!(
139                        clause_kind,
140                        MergeClauseKind::NotMatched | MergeClauseKind::NotMatchedByTarget
141                    ) {
142                        return parser_err!(
143                            format_args!("DELETE is not allowed in a {clause_kind} merge clause"),
144                            self.get_current_token().span.start
145                        );
146                    };
147
148                    let delete_token = self.get_current_token().clone();
149                    MergeAction::Delete {
150                        delete_token: delete_token.into(),
151                    }
152                }
153                Some(Keyword::INSERT) => {
154                    if !matches!(
155                        clause_kind,
156                        MergeClauseKind::NotMatched | MergeClauseKind::NotMatchedByTarget
157                    ) {
158                        return parser_err!(
159                            format_args!("INSERT is not allowed in a {clause_kind} merge clause"),
160                            self.get_current_token().span.start
161                        );
162                    };
163
164                    let insert_token = self.get_current_token().clone();
165                    let is_mysql = dialect_of!(self is MySqlDialect);
166
167                    let columns = self.parse_merge_clause_insert_columns(is_mysql)?;
168                    let (kind, kind_token) = if dialect_of!(self is BigQueryDialect | GenericDialect)
169                        && self.parse_keyword(Keyword::ROW)
170                    {
171                        (MergeInsertKind::Row, self.get_current_token().clone())
172                    } else {
173                        self.expect_keyword_is(Keyword::VALUES)?;
174                        let values_token = self.get_current_token().clone();
175                        let values = self.parse_values(is_mysql, false)?;
176                        (MergeInsertKind::Values(values), values_token)
177                    };
178                    let insert_predicate = if self.parse_keyword(Keyword::WHERE) {
179                        Some(self.parse_expr()?)
180                    } else {
181                        None
182                    };
183
184                    MergeAction::Insert(MergeInsertExpr {
185                        insert_token: insert_token.into(),
186                        columns,
187                        kind_token: kind_token.into(),
188                        kind,
189                        insert_predicate,
190                    })
191                }
192                _ => {
193                    return parser_err!(
194                        "expected UPDATE, DELETE or INSERT in merge clause",
195                        self.peek_token_ref().span.start
196                    );
197                }
198            };
199            clauses.push(MergeClause {
200                when_token: when_token.into(),
201                clause_kind,
202                predicate,
203                action: merge_clause,
204            });
205        }
206        Ok(clauses)
207    }
208
209    fn parse_merge_clause_insert_columns(
210        &mut self,
211        allow_empty: bool,
212    ) -> Result<Vec<ObjectName>, ParserError> {
213        self.parse_parenthesized_qualified_column_list(IsOptional::Optional, allow_empty)
214    }
215
216    fn parse_output(
217        &mut self,
218        start_keyword: Keyword,
219        start_token: TokenWithSpan,
220    ) -> Result<OutputClause, ParserError> {
221        let select_items = self.parse_projection()?;
222        let into_table = if start_keyword == Keyword::OUTPUT && self.peek_keyword(Keyword::INTO) {
223            self.expect_keyword_is(Keyword::INTO)?;
224            Some(self.parse_select_into()?)
225        } else {
226            None
227        };
228
229        Ok(if start_keyword == Keyword::OUTPUT {
230            OutputClause::Output {
231                output_token: start_token.into(),
232                select_items,
233                into_table,
234            }
235        } else {
236            OutputClause::Returning {
237                returning_token: start_token.into(),
238                select_items,
239            }
240        })
241    }
242}