qusql_parse/
with_query.rs

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

qusql_parse/
with_query.rs
