Skip to main content

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, vec, vec::Vec};
17
18use crate::{
19    ast::{
20        Merge, MergeAction, MergeClause, MergeClauseKind, MergeInsertExpr, MergeInsertKind,
21        MergeUpdateExpr, MergeUpdateKind, ObjectName, OutputClause, SetExpr,
22    },
23    dialect::{BigQueryDialect, GenericDialect, MySqlDialect},
24    keywords::Keyword,
25    parser::IsOptional,
26    tokenizer::Token,
27    tokenizer::TokenWithSpan,
28};
29
30use super::{Parser, ParserError};
31
32impl Parser<'_> {
33    /// Parse a `MERGE` statement, returning a `Box`ed SetExpr
34    ///
35    /// This is used to reduce the size of the stack frames in debug builds
36    pub(super) fn parse_merge_setexpr_boxed(
37        &mut self,
38        merge_token: TokenWithSpan,
39    ) -> Result<Box<SetExpr>, ParserError> {
40        Ok(Box::new(SetExpr::Merge(
41            self.parse_merge(merge_token)?.into(),
42        )))
43    }
44
45    /// Parse a `MERGE` statement
46    pub fn parse_merge(&mut self, merge_token: TokenWithSpan) -> Result<Merge, ParserError> {
47        let optimizer_hints = self.maybe_parse_optimizer_hints()?;
48        let into = self.parse_keyword(Keyword::INTO);
49
50        let table = self.parse_table_factor()?;
51
52        self.expect_keyword_is(Keyword::USING)?;
53        let source = self.parse_table_factor()?;
54        self.expect_keyword_is(Keyword::ON)?;
55        let on = self.parse_expr()?;
56        let clauses = self.parse_merge_clauses()?;
57        let output = match self.parse_one_of_keywords(&[Keyword::OUTPUT, Keyword::RETURNING]) {
58            Some(keyword) => Some(self.parse_output(keyword, self.get_current_token().clone())?),
59            None => None,
60        };
61
62        Ok(Merge {
63            merge_token: merge_token.into(),
64            optimizer_hints,
65            into,
66            table,
67            source,
68            on: Box::new(on),
69            clauses,
70            output,
71        })
72    }
73
74    fn parse_merge_clauses(&mut self) -> Result<Vec<MergeClause>, ParserError> {
75        let mut clauses = vec![];
76        loop {
77            if !(self.parse_keyword(Keyword::WHEN)) {
78                break;
79            }
80            let when_token = self.get_current_token().clone();
81
82            let mut clause_kind = MergeClauseKind::Matched;
83            if self.parse_keyword(Keyword::NOT) {
84                clause_kind = MergeClauseKind::NotMatched;
85            }
86            self.expect_keyword_is(Keyword::MATCHED)?;
87
88            if matches!(clause_kind, MergeClauseKind::NotMatched)
89                && self.parse_keywords(&[Keyword::BY, Keyword::SOURCE])
90            {
91                clause_kind = MergeClauseKind::NotMatchedBySource;
92            } else if matches!(clause_kind, MergeClauseKind::NotMatched)
93                && self.parse_keywords(&[Keyword::BY, Keyword::TARGET])
94            {
95                clause_kind = MergeClauseKind::NotMatchedByTarget;
96            }
97
98            let predicate = if self.parse_keyword(Keyword::AND) {
99                Some(self.parse_expr()?)
100            } else {
101                None
102            };
103
104            self.expect_keyword_is(Keyword::THEN)?;
105
106            let merge_clause = match self.parse_one_of_keywords(&[
107                Keyword::UPDATE,
108                Keyword::INSERT,
109                Keyword::DELETE,
110            ]) {
111                Some(Keyword::UPDATE) => {
112                    if matches!(
113                        clause_kind,
114                        MergeClauseKind::NotMatched | MergeClauseKind::NotMatchedByTarget
115                    ) {
116                        return parser_err!(
117                            format_args!("UPDATE is not allowed in a {clause_kind} merge clause"),
118                            self.get_current_token().span.start
119                        );
120                    }
121
122                    let update_token = self.get_current_token().clone();
123                    self.expect_keyword_is(Keyword::SET)?;
124                    let kind = if self.consume_token(&Token::Mul) {
125                        MergeUpdateKind::Wildcard
126                    } else {
127                        MergeUpdateKind::Set(self.parse_comma_separated(Parser::parse_assignment)?)
128                    };
129                    let update_predicate = if self.parse_keyword(Keyword::WHERE) {
130                        Some(self.parse_expr()?)
131                    } else {
132                        None
133                    };
134                    let delete_predicate = if self.parse_keyword(Keyword::DELETE) {
135                        let _ = self.expect_keyword(Keyword::WHERE)?;
136                        Some(self.parse_expr()?)
137                    } else {
138                        None
139                    };
140                    MergeAction::Update(MergeUpdateExpr {
141                        update_token: update_token.into(),
142                        kind,
143                        update_predicate,
144                        delete_predicate,
145                    })
146                }
147                Some(Keyword::DELETE) => {
148                    if matches!(
149                        clause_kind,
150                        MergeClauseKind::NotMatched | MergeClauseKind::NotMatchedByTarget
151                    ) {
152                        return parser_err!(
153                            format_args!("DELETE is not allowed in a {clause_kind} merge clause"),
154                            self.get_current_token().span.start
155                        );
156                    };
157
158                    let delete_token = self.get_current_token().clone();
159                    MergeAction::Delete {
160                        delete_token: delete_token.into(),
161                    }
162                }
163                Some(Keyword::INSERT) => {
164                    if !matches!(
165                        clause_kind,
166                        MergeClauseKind::NotMatched | MergeClauseKind::NotMatchedByTarget
167                    ) {
168                        return parser_err!(
169                            format_args!("INSERT is not allowed in a {clause_kind} merge clause"),
170                            self.get_current_token().span.start
171                        );
172                    };
173
174                    let insert_token = self.get_current_token().clone();
175
176                    if self.consume_token(&Token::Mul) {
177                        let star_token = self.get_current_token().clone();
178                        MergeAction::Insert(MergeInsertExpr {
179                            insert_token: insert_token.into(),
180                            columns: vec![],
181                            kind_token: star_token.into(),
182                            kind: MergeInsertKind::Wildcard,
183                            insert_predicate: None,
184                        })
185                    } else {
186                        let is_mysql = dialect_of!(self is MySqlDialect);
187                        let columns = self.parse_merge_clause_insert_columns(is_mysql)?;
188                        let (kind, kind_token) = if dialect_of!(self is BigQueryDialect | GenericDialect)
189                            && self.parse_keyword(Keyword::ROW)
190                        {
191                            (MergeInsertKind::Row, self.get_current_token().clone())
192                        } else {
193                            self.expect_keyword_is(Keyword::VALUES)?;
194                            let values_token = self.get_current_token().clone();
195                            let values = self.parse_values(is_mysql, false)?;
196                            (MergeInsertKind::Values(values), values_token)
197                        };
198                        let insert_predicate = if self.parse_keyword(Keyword::WHERE) {
199                            Some(self.parse_expr()?)
200                        } else {
201                            None
202                        };
203
204                        MergeAction::Insert(MergeInsertExpr {
205                            insert_token: insert_token.into(),
206                            columns,
207                            kind_token: kind_token.into(),
208                            kind,
209                            insert_predicate,
210                        })
211                    }
212                }
213                _ => {
214                    return parser_err!(
215                        "expected UPDATE, DELETE or INSERT in merge clause",
216                        self.peek_token_ref().span.start
217                    );
218                }
219            };
220            clauses.push(MergeClause {
221                when_token: when_token.into(),
222                clause_kind,
223                predicate,
224                action: merge_clause,
225            });
226        }
227        Ok(clauses)
228    }
229
230    fn parse_merge_clause_insert_columns(
231        &mut self,
232        allow_empty: bool,
233    ) -> Result<Vec<ObjectName>, ParserError> {
234        self.parse_parenthesized_qualified_column_list(IsOptional::Optional, allow_empty)
235    }
236
237    /// Parses an `OUTPUT` clause if present (MSSQL).
238    pub(super) fn maybe_parse_output_clause(
239        &mut self,
240    ) -> Result<Option<OutputClause>, ParserError> {
241        if self.parse_keyword(Keyword::OUTPUT) {
242            Ok(Some(self.parse_output(
243                Keyword::OUTPUT,
244                self.get_current_token().clone(),
245            )?))
246        } else {
247            Ok(None)
248        }
249    }
250
251    pub(super) fn parse_output(
252        &mut self,
253        start_keyword: Keyword,
254        start_token: TokenWithSpan,
255    ) -> Result<OutputClause, ParserError> {
256        let select_items = self.parse_projection()?;
257        let into_table = if start_keyword == Keyword::OUTPUT && self.peek_keyword(Keyword::INTO) {
258            self.expect_keyword_is(Keyword::INTO)?;
259            Some(self.parse_select_into()?)
260        } else {
261            None
262        };
263
264        Ok(if start_keyword == Keyword::OUTPUT {
265            OutputClause::Output {
266                output_token: start_token.into(),
267                select_items,
268                into_table,
269            }
270        } else {
271            OutputClause::Returning {
272                returning_token: start_token.into(),
273                select_items,
274            }
275        })
276    }
277}