Skip to main content

qraft_core/
cte.rs

1use std::marker::PhantomData;
2
3use crate::{
4    Query,
5    builder::{delete::Delete, update::UpdateFrom},
6    query::{LowerProject, Table, WithSelect, select},
7};
8
9/// A lowered common table expression attached to a statement.
10#[derive(Debug, Clone)]
11pub struct Cte {
12    pub(crate) name: &'static str,
13    pub(crate) columns: Vec<&'static str>,
14    pub(crate) query: Query,
15    pub(crate) recursive: bool,
16}
17
18/// A typed CTE definition before recursive mode has been selected.
19#[derive(Debug, Clone)]
20pub struct CteDefinition<T> {
21    name: &'static str,
22    columns: &'static [&'static str],
23    query: Query,
24    marker: PhantomData<T>,
25}
26
27/// Converts one or more typed CTE definitions into lowered CTEs.
28pub trait IntoCtes {
29    #[doc(hidden)]
30    fn into_ctes(self, recursive: bool) -> Vec<Cte>;
31}
32
33/// A top-level `with (...)` wrapper that can start statements.
34#[derive(Debug, Clone)]
35pub struct WithClause {
36    ctes: Vec<Cte>,
37}
38
39impl<T> CteDefinition<T> {
40    /// Creates a typed CTE definition from a derived schema module.
41    pub fn new<Q>(name: &'static str, columns: &'static [&'static str], query: Q) -> Self
42    where
43        Q: Into<Query>,
44    {
45        Self {
46            name,
47            columns,
48            query: query.into(),
49            marker: PhantomData,
50        }
51    }
52
53    fn into_cte(self, recursive: bool) -> Cte {
54        Cte {
55            name: self.name,
56            columns: self.columns.to_vec(),
57            query: self.query,
58            recursive,
59        }
60    }
61}
62
63impl<T> IntoCtes for CteDefinition<T> {
64    fn into_ctes(self, recursive: bool) -> Vec<Cte> {
65        vec![self.into_cte(recursive)]
66    }
67}
68
69impl IntoCtes for Cte {
70    fn into_ctes(mut self, recursive: bool) -> Vec<Cte> {
71        self.recursive |= recursive;
72        vec![self]
73    }
74}
75
76impl IntoCtes for Vec<Cte> {
77    fn into_ctes(mut self, recursive: bool) -> Vec<Cte> {
78        if recursive {
79            for cte in &mut self {
80                cte.recursive = true;
81            }
82        }
83        self
84    }
85}
86
87impl<T> IntoCtes for Vec<CteDefinition<T>> {
88    fn into_ctes(self, recursive: bool) -> Vec<Cte> {
89        self.into_iter()
90            .map(|cte| cte.into_cte(recursive))
91            .collect()
92    }
93}
94
95macro_rules! impl_into_ctes_tuple {
96    ($($T:ident),+) => {
97        impl<$($T,)+> IntoCtes for ($($T,)+)
98        where
99            $($T: IntoCtes,)+
100        {
101            fn into_ctes(self, recursive: bool) -> Vec<Cte> {
102                #[allow(non_snake_case)]
103                let ($($T,)+) = self;
104                let mut out = Vec::new();
105                $(
106                    out.extend($T.into_ctes(recursive));
107                )+
108                out
109            }
110        }
111    };
112}
113
114crate::impl_for_all_tuples!(impl_into_ctes_tuple);
115
116/// Starts a `with (...)` wrapper from one or more typed CTE definitions.
117pub fn with<C>(ctes: C) -> WithClause
118where
119    C: IntoCtes,
120{
121    WithClause {
122        ctes: ctes.into_ctes(false),
123    }
124}
125
126/// Starts a `with recursive (...)` wrapper from one or more typed CTE definitions.
127pub fn with_recursive<C>(ctes: C) -> WithClause
128where
129    C: IntoCtes,
130{
131    WithClause {
132        ctes: ctes.into_ctes(true),
133    }
134}
135
136impl WithClause {
137    /// Starts a select statement with the attached CTEs.
138    pub fn select<P>(self, project: P) -> WithSelect<P>
139    where
140        P: LowerProject,
141    {
142        let mut ctes = self.ctes.into_iter();
143        let first = ctes.next().expect("with() requires at least one CTE");
144        let mut stmt = select(project).with(first);
145        for cte in ctes {
146            stmt = stmt.with(cte);
147        }
148        stmt
149    }
150
151    /// Starts an update statement with the attached CTEs.
152    pub fn update<M, T>(self, table: T) -> UpdateFrom<M>
153    where
154        T: crate::builder::update::UpdateTable<M>,
155    {
156        let mut stmt = crate::builder::update::update(table);
157        for cte in self.ctes {
158            stmt = stmt.with(cte);
159        }
160        stmt
161    }
162
163    /// Starts a delete statement with the attached CTEs.
164    pub fn delete<T>(self, table: Table<T>) -> Delete<T> {
165        let mut stmt = crate::builder::delete::delete_from(table);
166        for cte in self.ctes {
167            stmt = stmt.with(cte);
168        }
169        stmt
170    }
171}