Skip to main content

shape_ast/parser/queries/
with.rs

1//! WITH query (CTE) parsing implementation
2
3use crate::ast::{Cte, Query, WithQuery};
4use crate::error::{Result, ShapeError};
5use crate::parser::{Rule, pair_location};
6use pest::iterators::Pair;
7
8use super::alert::parse_alert_query;
9
10/// Parse an inner query (used in CTEs and as the main query)
11fn parse_inner_query(pair: Pair<Rule>) -> Result<Query> {
12    let inner = pair
13        .into_inner()
14        .next()
15        .ok_or_else(|| ShapeError::ParseError {
16            message: "expected query in CTE".to_string(),
17            location: None,
18        })?;
19
20    match inner.as_rule() {
21        Rule::alert_query => Ok(Query::Alert(parse_alert_query(inner)?)),
22        Rule::with_query => Ok(Query::With(parse_with_query(inner)?)),
23        _ => Err(ShapeError::ParseError {
24            message: format!("unexpected query type in CTE: {:?}", inner.as_rule()),
25            location: Some(pair_location(&inner)),
26        }),
27    }
28}
29
30/// Parse a single CTE definition
31fn parse_cte_def(pair: Pair<Rule>) -> Result<Cte> {
32    let inner = pair.into_inner();
33    let mut recursive = false;
34    let mut name = String::new();
35    let mut columns = None;
36    let mut query = None;
37
38    for item in inner {
39        match item.as_rule() {
40            Rule::recursive_keyword => {
41                recursive = true;
42            }
43            Rule::ident => {
44                name = item.as_str().to_string();
45            }
46            Rule::cte_columns => {
47                let cols: Vec<String> = item
48                    .into_inner()
49                    .filter(|p| p.as_rule() == Rule::ident)
50                    .map(|p| p.as_str().to_string())
51                    .collect();
52                columns = Some(cols);
53            }
54            Rule::inner_query => {
55                query = Some(parse_inner_query(item)?);
56            }
57            _ => {}
58        }
59    }
60
61    let query = query.ok_or_else(|| ShapeError::ParseError {
62        message: "CTE definition missing query".to_string(),
63        location: None,
64    })?;
65
66    Ok(Cte {
67        name,
68        columns,
69        query: Box::new(query),
70        recursive,
71    })
72}
73
74/// Parse a list of CTEs
75fn parse_cte_list(pair: Pair<Rule>) -> Result<Vec<Cte>> {
76    pair.into_inner()
77        .filter(|p| p.as_rule() == Rule::cte_def)
78        .map(parse_cte_def)
79        .collect()
80}
81
82/// Parse a WITH query (query with Common Table Expressions)
83///
84/// Syntax:
85/// ```text
86/// WITH
87///     cte1 AS (query1),
88///     cte2 AS (query2),
89///     RECURSIVE cte3(col1, col2) AS (query3)
90/// main_query
91/// ```
92pub fn parse_with_query(pair: Pair<Rule>) -> Result<WithQuery> {
93    let inner = pair.into_inner();
94    let mut ctes = Vec::new();
95    let mut main_query = None;
96
97    for item in inner {
98        match item.as_rule() {
99            Rule::cte_list => {
100                ctes = parse_cte_list(item)?;
101            }
102            Rule::inner_query => {
103                main_query = Some(parse_inner_query(item)?);
104            }
105            _ => {}
106        }
107    }
108
109    let query = main_query.ok_or_else(|| ShapeError::ParseError {
110        message: "WITH clause missing main query".to_string(),
111        location: None,
112    })?;
113
114    Ok(WithQuery {
115        ctes,
116        query: Box::new(query),
117    })
118}