Skip to main content

keelson_core/clause/
with.rs

1use crate::writer::{Expression, SqlWriter};
2
3use super::cte::Cte;
4use super::write_present;
5
6/// `WITH [RECURSIVE] <cte>, <cte>, …`
7///
8/// From PostgreSQL 17: `WITH [ RECURSIVE ] with_query [, ...]`. `RECURSIVE` is a
9/// property of the whole list rather than of one entry — it makes every name in the
10/// list visible to every entry, which is also what lets two CTEs refer to each
11/// other.
12#[derive(Debug, Clone, Default)]
13pub struct With {
14    /// Whether the list is `WITH RECURSIVE`.
15    pub recursive: bool,
16    /// The common table expressions, in order.
17    pub ctes: Vec<Cte>,
18}
19
20impl With {
21    /// Append one CTE.
22    pub fn append_cte(&mut self, cte: Cte) {
23        self.ctes.push(cte);
24    }
25
26    /// Make the list recursive.
27    pub fn set_recursive(&mut self, recursive: bool) {
28        self.recursive = recursive;
29    }
30
31    /// Whether the clause is absent.
32    pub fn is_empty(&self) -> bool {
33        self.ctes.is_empty()
34    }
35}
36
37impl Expression for With {
38    fn write_sql(&self, w: &mut SqlWriter<'_>) {
39        // `WITH RECURSIVE` with no CTEs is not a clause, so the list gates the
40        // keyword — including the RECURSIVE that a mod may have set first.
41        let prefix = if self.recursive {
42            "WITH RECURSIVE "
43        } else {
44            "WITH "
45        };
46        write_present(w, &self.ctes, prefix, ", ", "");
47    }
48}
49
50/// A statement that accepts common table expressions.
51pub trait HasWith {
52    /// The `WITH` clause to modify.
53    fn with_mut(&mut self) -> &mut With;
54}
55
56impl HasWith for With {
57    fn with_mut(&mut self) -> &mut With {
58        self
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use keelson_sqlcheck::testing::assert_frag_sql;
65
66    use super::*;
67    use crate::dialect::testing::Numbered;
68    use crate::expr::{Expr, arg};
69    use crate::value::Value;
70    use crate::writer::build;
71
72    /// A `WITH` prefixes a statement, so the frame follows it rather than
73    /// surrounding it. Two frames: one whose body reads the CTEs, and one for the
74    /// cases where the clause writes nothing and there is no CTE to read.
75    const FRAME: &str = r#"{} SELECT * FROM "a""#;
76    const EMPTY_FRAME: &str = r#"{} SELECT "id" FROM users"#;
77
78    /// A real sub-query: `SELECT "id" FROM posts WHERE "id" = $n`. The placeholder
79    /// is compared against a column so PostgreSQL can give it a type.
80    fn sub(v: i32) -> Expr {
81        Expr::join((Expr::raw(r#"SELECT "id" FROM posts WHERE "id" ="#), arg(v)))
82    }
83
84    fn sub_sql(n: usize) -> String {
85        format!(r#"SELECT "id" FROM posts WHERE "id" = ${n}"#)
86    }
87
88    fn sql(w: &With) -> String {
89        build(&Numbered, w).expect("render").0
90    }
91
92    #[test]
93    fn an_empty_with_writes_nothing_not_even_the_keyword() {
94        assert_frag_sql(EMPTY_FRAME, &sql(&With::default()), "");
95        assert!(With::default().is_empty());
96    }
97
98    #[test]
99    fn recursive_alone_is_still_an_absent_clause() {
100        let mut with = With::default();
101        with.set_recursive(true);
102        assert_frag_sql(EMPTY_FRAME, &sql(&with), "");
103    }
104
105    #[test]
106    fn ctes_are_comma_separated_and_share_one_placeholder_run() {
107        let mut with = With::default();
108        with.append_cte(Cte::new("a", sub(1)));
109        with.append_cte(Cte::new("b", sub(2)));
110
111        let (rendered, args) = build(&Numbered, &with).unwrap();
112        assert_frag_sql(
113            FRAME,
114            &rendered,
115            &format!(r#"WITH "a" AS ({}), "b" AS ({})"#, sub_sql(1), sub_sql(2)),
116        );
117        assert_eq!(args, vec![Value::I32(1), Value::I32(2)]);
118
119        // RECURSIVE is a property of the whole list, not of one CTE, and a list
120        // that happens not to recurse is still allowed to carry it.
121        with.set_recursive(true);
122        assert_frag_sql(
123            FRAME,
124            &sql(&with),
125            &format!(
126                r#"WITH RECURSIVE "a" AS ({}), "b" AS ({})"#,
127                sub_sql(1),
128                sub_sql(2)
129            ),
130        );
131    }
132
133    #[test]
134    fn an_absent_cte_takes_its_comma_with_it() {
135        // `WITH "a" AS (…), ` does not parse, so a separator is written only
136        // between CTEs that are both actually there.
137        let mut with = With::default();
138        with.append_cte(Cte::default());
139        assert_frag_sql(EMPTY_FRAME, &sql(&with), "");
140
141        with.append_cte(Cte::new("a", sub(1)));
142        with.append_cte(Cte::default());
143        assert_frag_sql(
144            FRAME,
145            &sql(&with),
146            &format!(r#"WITH "a" AS ({})"#, sub_sql(1)),
147        );
148    }
149}