Skip to main content

sqlparser/parser/
procedural.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use super::*;
19
20impl<'a> Parser<'a> {
21    pub(super) fn parse_do_statement(&mut self) -> Result<DoStatement, ParserError> {
22        let language_before = if self.parse_keyword(Keyword::LANGUAGE) {
23            Some(self.parse_identifier()?)
24        } else {
25            None
26        };
27        let body = self.parse_value()?;
28        if !matches!(
29            body.value,
30            Value::SingleQuotedString(_)
31                | Value::EscapedStringLiteral(_)
32                | Value::UnicodeStringLiteral(_)
33                | Value::DollarQuotedString(_)
34        ) {
35            return self.expected_ref("a string literal", self.get_current_token());
36        }
37        let language_after = if self.parse_keyword(Keyword::LANGUAGE) {
38            Some(self.parse_identifier()?)
39        } else {
40            None
41        };
42        if language_before.is_some() && language_after.is_some() {
43            return Err(ParserError::ParserError(
44                "LANGUAGE specified more than once".into(),
45            ));
46        }
47        Ok(DoStatement {
48            body,
49            language: language_before.clone().or(language_after),
50            language_before_body: language_before.is_some(),
51        })
52    }
53
54    /// Parse one complete PostgreSQL PL/pgSQL block.
55    pub fn parse_plpgsql(&mut self) -> Result<PlPgSqlBlock, ParserError> {
56        let declarations = if self.parse_keyword(Keyword::DECLARE) {
57            let mut declarations = Vec::new();
58            while !self.peek_keyword(Keyword::BEGIN) {
59                if matches!(self.peek_token_ref().token, Token::EOF) {
60                    return self.expected_ref("BEGIN", self.peek_token_ref());
61                }
62                let name = self.parse_identifier()?;
63                let data_type = self.parse_data_type()?;
64                let initializer = if self.consume_token(&Token::Assignment)
65                    || self.consume_token(&Token::Eq)
66                    || self.parse_keyword(Keyword::DEFAULT)
67                {
68                    Some(self.parse_expr()?)
69                } else {
70                    None
71                };
72                self.expect_token(&Token::SemiColon)?;
73                declarations.push(PlPgSqlDeclaration {
74                    name,
75                    data_type,
76                    initializer,
77                });
78            }
79            declarations
80        } else {
81            Vec::new()
82        };
83        self.expect_keyword_is(Keyword::BEGIN)?;
84        let statements = self.parse_plpgsql_statements(&[Keyword::END])?;
85        self.expect_keyword_is(Keyword::END)?;
86        let _ = self.consume_token(&Token::SemiColon);
87        self.expect_token(&Token::EOF)?;
88        Ok(PlPgSqlBlock {
89            declarations,
90            statements,
91        })
92    }
93
94    fn parse_plpgsql_statements(
95        &mut self,
96        terminal_keywords: &[Keyword],
97    ) -> Result<Vec<PlPgSqlStatement>, ParserError> {
98        let _guard = self.recursion_counter.try_decrease()?;
99        let mut statements = Vec::new();
100        while !terminal_keywords
101            .iter()
102            .any(|keyword| self.peek_keyword(*keyword))
103        {
104            if matches!(self.peek_token_ref().token, Token::EOF) {
105                return self.expected_ref("a PL/pgSQL statement", self.peek_token_ref());
106            }
107            statements.push(self.parse_plpgsql_statement()?);
108            self.expect_token(&Token::SemiColon)?;
109        }
110        Ok(statements)
111    }
112
113    fn parse_plpgsql_statement(&mut self) -> Result<PlPgSqlStatement, ParserError> {
114        if self.parse_keyword(Keyword::IF) {
115            return self.parse_plpgsql_if();
116        }
117        if self.parse_keyword(Keyword::RETURN) {
118            return self.parse_expr().map(PlPgSqlStatement::Return);
119        }
120        if self.parse_keyword(Keyword::GET) {
121            self.expect_keyword_is(Keyword::DIAGNOSTICS)?;
122            let target = self.parse_identifier()?;
123            self.expect_token(&Token::Eq)?;
124            self.expect_keyword_is(Keyword::ROW_COUNT)?;
125            return Ok(PlPgSqlStatement::GetRowCount { target });
126        }
127        if self.parse_keyword(Keyword::RAISE) {
128            self.expect_keyword_is(Keyword::EXCEPTION)?;
129            let format = self.parse_value()?;
130            if !matches!(
131                format.value,
132                Value::SingleQuotedString(_)
133                    | Value::EscapedStringLiteral(_)
134                    | Value::UnicodeStringLiteral(_)
135                    | Value::DollarQuotedString(_)
136            ) {
137                return self.expected_ref("a string literal", self.get_current_token());
138            }
139            let mut arguments = Vec::new();
140            while self.consume_token(&Token::Comma) {
141                arguments.push(self.parse_expr()?);
142            }
143            let hint = if self.parse_keywords(&[Keyword::USING, Keyword::HINT]) {
144                self.expect_token(&Token::Eq)?;
145                Some(self.parse_expr()?)
146            } else {
147                None
148            };
149            return Ok(PlPgSqlStatement::RaiseException {
150                format,
151                arguments,
152                hint,
153            });
154        }
155        if matches!(
156            &self.peek_token_ref().token,
157            Token::Word(word)
158                if matches!(
159                    word.keyword,
160                    Keyword::SELECT
161                        | Keyword::WITH
162                        | Keyword::VALUES
163                        | Keyword::INSERT
164                        | Keyword::UPDATE
165                        | Keyword::DELETE
166                )
167        ) {
168            return self
169                .parse_statement()
170                .map(Box::new)
171                .map(PlPgSqlStatement::Sql);
172        }
173
174        let target = self.parse_object_name(false)?;
175        if !self.consume_token(&Token::Assignment) {
176            self.expect_token(&Token::Eq)?;
177        }
178        let value = self.parse_expr()?;
179        Ok(PlPgSqlStatement::Assignment { target, value })
180    }
181
182    fn parse_plpgsql_if(&mut self) -> Result<PlPgSqlStatement, ParserError> {
183        let mut branches = Vec::new();
184        let condition = self.parse_expr()?;
185        self.expect_keyword_is(Keyword::THEN)?;
186        let statements = self.parse_plpgsql_statements(&[
187            Keyword::ELSIF,
188            Keyword::ELSEIF,
189            Keyword::ELSE,
190            Keyword::END,
191        ])?;
192        branches.push(PlPgSqlIfBranch {
193            condition,
194            statements,
195        });
196        while self.parse_keyword(Keyword::ELSIF) || self.parse_keyword(Keyword::ELSEIF) {
197            let condition = self.parse_expr()?;
198            self.expect_keyword_is(Keyword::THEN)?;
199            let statements = self.parse_plpgsql_statements(&[
200                Keyword::ELSIF,
201                Keyword::ELSEIF,
202                Keyword::ELSE,
203                Keyword::END,
204            ])?;
205            branches.push(PlPgSqlIfBranch {
206                condition,
207                statements,
208            });
209        }
210        let else_statements = if self.parse_keyword(Keyword::ELSE) {
211            Some(self.parse_plpgsql_statements(&[Keyword::END])?)
212        } else {
213            None
214        };
215        self.expect_keywords(&[Keyword::END, Keyword::IF])?;
216        Ok(PlPgSqlStatement::If {
217            branches,
218            else_statements,
219        })
220    }
221}