sql_parse/
with_query.rs

1use alloc::{boxed::Box, vec::Vec};
2
3use crate::{
4    keywords::Keyword,
5    lexer::Token,
6    parser::{ParseError, Parser},
7    statement::parse_statement,
8    Identifier, Span, Spanned, Statement,
9};
10
11#[derive(Clone, Debug)]
12pub struct WithBlock<'a> {
13    // Identifier for the with block
14    pub identifier: Identifier<'a>,
15    // Span of AS
16    pub as_span: Span,
17    // Span of (
18    pub lparen_span: Span,
19    // The statement within the with block, will be one of select, update, insert or delete
20    pub statement: Statement<'a>,
21    // Span of )
22    pub rparen_span: Span,
23}
24
25impl<'a> Spanned for WithBlock<'a> {
26    fn span(&self) -> Span {
27        self.identifier
28            .span()
29            .join_span(&self.as_span)
30            .join_span(&self.lparen_span)
31            .join_span(&self.statement)
32            .join_span(&self.rparen_span)
33    }
34}
35
36/// Represent a with query statement
37/// ```
38/// # use sql_parse::{SQLDialect, SQLArguments, ParseOptions, parse_statements, WithQuery, Statement, Issues};
39/// # let options = ParseOptions::new().dialect(SQLDialect::PostgreSQL);
40/// #
41/// let sql = "WITH ids AS (DELETE FROM things WHERE number=42) INSERT INTO deleted (id) SELECT id FROM ids;";
42/// let mut issues = Issues::new(sql);
43/// let mut stmts = parse_statements(sql, &mut issues, &options);
44///
45/// # assert!(issues.is_ok());
46/// #
47/// let delete: WithQuery = match stmts.pop() {
48///     Some(Statement::WithQuery(d)) => d,
49///     _ => panic!("We should get a with statement")
50/// };
51///
52/// assert!(delete.with_blocks[0].identifier.as_str() == "ids");
53/// ```
54#[derive(Clone, Debug)]
55pub struct WithQuery<'a> {
56    // Span of WITH
57    pub with_span: Span,
58    // The comma seperated with blocks
59    pub with_blocks: Vec<WithBlock<'a>>,
60    // The final statement of the with query, will be one of select, update, insert, delete or merge
61    pub statement: Box<Statement<'a>>,
62}
63
64impl<'a> Spanned for WithQuery<'a> {
65    fn span(&self) -> Span {
66        self.with_span
67            .join_span(&self.with_blocks)
68            .join_span(&self.statement)
69    }
70}
71
72pub(crate) fn parse_with_query<'a>(
73    parser: &mut Parser<'a, '_>,
74) -> Result<WithQuery<'a>, ParseError> {
75    let with_span = parser.consume_keyword(Keyword::WITH)?;
76    let mut with_blocks = Vec::new();
77    loop {
78        let identifier = parser.consume_plain_identifier()?;
79        let as_span = parser.consume_keyword(Keyword::AS)?;
80        let lparen_span = parser.consume_token(Token::LParen)?;
81        let statement =
82            parser.recovered(
83                "')'",
84                &|t| t == &Token::RParen,
85                |parser| match parse_statement(parser)? {
86                    Some(v) => Ok(Some(v)),
87                    None => {
88                        parser.expected_error("Statement");
89                        Ok(None)
90                    }
91                },
92            )?;
93        let rparen_span = parser.consume_token(Token::RParen)?;
94        let statement = match statement {
95            Some(v) => {
96                if !matches!(
97                    &v,
98                    Statement::Select(_)
99                        | Statement::InsertReplace(_)
100                        | Statement::Update(_)
101                        | Statement::Delete(_)
102                ) {
103                    parser.err(
104                        "Only SELECT, INSERT, UPDATE or DELETE allowed within WITH query",
105                        &v.span(),
106                    );
107                }
108                v
109            }
110            None => Statement::Begin(lparen_span.clone()),
111        };
112        with_blocks.push(WithBlock {
113            identifier,
114            as_span,
115            lparen_span,
116            statement,
117            rparen_span,
118        });
119        if parser.skip_token(Token::Comma).is_none() {
120            break;
121        }
122    }
123    let statement = match parse_statement(parser)? {
124        Some(v) => {
125            // TODO merge statements are also allowed here
126            if !matches!(
127                &v,
128                Statement::Select(_)
129                    | Statement::InsertReplace(_)
130                    | Statement::Update(_)
131                    | Statement::Delete(_)
132            ) {
133                parser.err(
134                    "Only SELECT, INSERT, UPDATE or DELETE allowed as WITH query",
135                    &v.span(),
136                );
137            }
138            Box::new(v)
139        }
140        None => parser.expected_failure("Statement")?,
141    };
142    let res = WithQuery {
143        with_span,
144        with_blocks,
145        statement,
146    };
147    if !parser.options.dialect.is_postgresql() {
148        parser.err("WITH statements only supported by postgresql", &res.span());
149    }
150    Ok(res)
151}
sql_parse/with_query.rs

sql_parse/
with_query.rs
