Skip to main content

toasty_core/stmt/
query.rs

1use super::{
2    Delete, Expr, ExprSet, Limit, Node, OrderBy, Path, Returning, Select, Source, Statement,
3    Update, UpdateTarget, Values, Visit, VisitMut, With,
4};
5use crate::stmt::{self, Filter};
6
7/// A query statement that reads data from the database.
8///
9/// `Query` wraps a set expression body (typically a [`Select`]) with optional
10/// ordering, limits, CTEs, and row-level locks. It is the read-side counterpart
11/// to [`Insert`], [`Update`], and [`Delete`].
12///
13/// # Examples
14///
15/// ```ignore
16/// use toasty_core::stmt::{Query, Values, ExprSet};
17///
18/// // A unit query that returns one empty row
19/// let q = Query::unit();
20/// assert!(matches!(q.body, ExprSet::Values(_)));
21/// assert!(!q.single);
22/// ```
23#[derive(Debug, Clone, PartialEq)]
24pub struct Query {
25    /// Optional common table expressions (CTEs) for this query.
26    pub with: Option<With>,
27
28    /// The body of the query. Either `SELECT`, `UNION`, `VALUES`, or possibly
29    /// other types of queries depending on database support.
30    pub body: ExprSet,
31
32    /// When `true`, the query returns a single record instead of a list.
33    ///
34    /// This is semantically different from `LIMIT 1`: it indicates there can
35    /// only ever be one matching result. The return type becomes `Record`
36    /// instead of `List`.
37    pub single: bool,
38
39    /// Optional `ORDER BY` clause.
40    pub order_by: Option<OrderBy>,
41
42    /// Optional `LIMIT` and `OFFSET` clause.
43    pub limit: Option<Limit>,
44
45    /// Row-level locks (`FOR UPDATE`, `FOR SHARE`).
46    pub locks: Vec<Lock>,
47}
48
49/// A row-level lock to acquire when executing a query.
50///
51/// Corresponds to SQL's `FOR UPDATE` and `FOR SHARE` clauses. Only meaningful
52/// for SQL databases that support row-level locking.
53///
54/// # Examples
55///
56/// ```ignore
57/// use toasty_core::stmt::Lock;
58///
59/// let lock = Lock::Update;
60/// ```
61#[derive(Debug, Clone, PartialEq)]
62pub enum Lock {
63    /// `FOR UPDATE` -- acquire an exclusive lock on matched rows.
64    Update,
65    /// `FOR SHARE` -- acquire a shared lock on matched rows.
66    Share,
67}
68
69/// Builder for constructing [`Query`] instances with optional clauses.
70///
71/// Created via [`Query::builder`]. Allows chaining calls to add CTEs, filters,
72/// returning clauses, and locks before calling [`build`](QueryBuilder::build).
73///
74/// # Examples
75///
76/// ```ignore
77/// use toasty_core::stmt::{Query, Select, Source, Filter};
78///
79/// let select = Select::new(Source::from(ModelId(0)), Filter::ALL);
80/// let query = Query::builder(select).build();
81/// ```
82#[derive(Debug)]
83pub struct QueryBuilder {
84    query: Query,
85}
86
87impl Query {
88    /// Creates a new query with the given body and default options (no ordering,
89    /// no limit, not single, no locks).
90    pub fn new(body: impl Into<ExprSet>) -> Self {
91        Self {
92            with: None,
93            body: body.into(),
94            single: false,
95            order_by: None,
96            limit: None,
97            locks: vec![],
98        }
99    }
100
101    /// Creates a new query that returns exactly one record (`single = true`).
102    pub fn new_single(body: impl Into<ExprSet>) -> Self {
103        Self {
104            with: None,
105            body: body.into(),
106            single: true,
107            order_by: None,
108            limit: None,
109            locks: vec![],
110        }
111    }
112
113    /// Creates a new `SELECT` query from a source and filter.
114    pub fn new_select(source: impl Into<Source>, filter: impl Into<Filter>) -> Self {
115        Self::builder(Select::new(source, filter)).build()
116    }
117
118    /// Returns a [`QueryBuilder`] initialized with the given body.
119    pub fn builder(body: impl Into<ExprSet>) -> QueryBuilder {
120        QueryBuilder {
121            query: Query::new(body),
122        }
123    }
124
125    /// Creates a unit query that produces one empty row (empty `VALUES`).
126    pub fn unit() -> Self {
127        Query::new(Values::default())
128    }
129
130    /// Creates a query whose body is a `VALUES` expression.
131    pub fn values(values: impl Into<Values>) -> Self {
132        Self {
133            with: None,
134            body: ExprSet::Values(values.into()),
135            single: false,
136            order_by: None,
137            limit: None,
138            locks: vec![],
139        }
140    }
141
142    /// Converts this query into an [`Update`] statement targeting the same
143    /// source. The query must have a `SELECT` body with a model source.
144    pub fn update(self) -> Update {
145        let ExprSet::Select(select) = &self.body else {
146            todo!("stmt={self:#?}");
147        };
148
149        assert!(select.source.is_model());
150
151        stmt::Update {
152            target: UpdateTarget::Query(Box::new(self)),
153            assignments: stmt::Assignments::default(),
154            filter: Filter::default(),
155            condition: stmt::Condition::default(),
156            returning: None,
157        }
158    }
159
160    /// Converts this query into a [`Delete`] statement. The query body must
161    /// be a `SELECT`.
162    pub fn delete(self) -> Delete {
163        match self.body {
164            ExprSet::Select(select) => Delete {
165                from: select.source,
166                filter: select.filter,
167                returning: None,
168                condition: Default::default(),
169            },
170            _ => todo!("{self:#?}"),
171        }
172    }
173
174    /// Adds a filter to this query's `SELECT` body.
175    ///
176    /// # Panics
177    ///
178    /// Panics if the query body is not a `SELECT`.
179    pub fn add_filter(&mut self, filter: impl Into<Filter>) {
180        self.body.as_select_mut_unwrap().add_filter(filter);
181    }
182
183    /// Adds an association include path to this query's `SELECT` body.
184    pub fn include(&mut self, path: impl Into<Path>) {
185        match &mut self.body {
186            ExprSet::Select(body) => body.include(path),
187            _ => todo!(),
188        }
189    }
190}
191
192impl Statement {
193    /// Returns `true` if this statement is a [`Query`].
194    pub fn is_query(&self) -> bool {
195        matches!(self, Statement::Query(_))
196    }
197
198    /// Attempts to return a reference to an inner [`Query`].
199    ///
200    /// * If `self` is a [`Statement::Query`], a reference to the inner [`Query`] is
201    ///   returned wrapped in [`Some`].
202    /// * Else, [`None`] is returned.
203    pub fn as_query(&self) -> Option<&Query> {
204        match self {
205            Self::Query(query) => Some(query),
206            _ => None,
207        }
208    }
209
210    /// Returns a mutable reference to the inner [`Query`], if this is a query statement.
211    ///
212    /// * If `self` is a [`Statement::Query`], a mutable reference to the inner [`Query`] is
213    ///   returned wrapped in [`Some`].
214    /// * Else, [`None`] is returned.
215    pub fn as_query_mut(&mut self) -> Option<&mut Query> {
216        match self {
217            Self::Query(query) => Some(query),
218            _ => None,
219        }
220    }
221
222    /// Consumes `self` and attempts to return the inner [`Query`].
223    ///
224    /// Returns `None` if `self` is not a [`Statement::Query`].
225    pub fn into_query(self) -> Option<Query> {
226        match self {
227            Self::Query(query) => Some(query),
228            _ => None,
229        }
230    }
231
232    /// Consumes `self` and returns the inner [`Query`].
233    ///
234    /// # Panics
235    ///
236    /// If `self` is not a [`Statement::Query`].
237    #[track_caller]
238    pub fn into_query_unwrap(self) -> Query {
239        match self {
240            Self::Query(query) => query,
241            v => panic!("expected `Query`, found {v:#?}"),
242        }
243    }
244}
245
246impl From<Query> for Statement {
247    fn from(value: Query) -> Self {
248        Self::Query(value)
249    }
250}
251
252impl Node for Query {
253    fn visit<V: Visit>(&self, mut visit: V) {
254        visit.visit_stmt_query(self);
255    }
256
257    fn visit_mut<V: VisitMut>(&mut self, mut visit: V) {
258        visit.visit_stmt_query_mut(self);
259    }
260}
261
262impl QueryBuilder {
263    /// Sets the `WITH` (CTE) clause for the query being built.
264    pub fn with(mut self, with: impl Into<With>) -> Self {
265        self.query.with = Some(with.into());
266        self
267    }
268
269    /// Sets the row-level locks for the query being built.
270    pub fn locks(mut self, locks: impl Into<Vec<Lock>>) -> Self {
271        self.query.locks = locks.into();
272        self
273    }
274
275    /// Sets the filter on the query's `SELECT` body.
276    ///
277    /// # Panics
278    ///
279    /// Panics if the query body is not a `SELECT`.
280    pub fn filter(mut self, filter: impl Into<Filter>) -> Self {
281        let filter = filter.into();
282
283        match &mut self.query.body {
284            ExprSet::Select(select) => {
285                select.filter = filter;
286            }
287            _ => todo!("query={self:#?}"),
288        }
289
290        self
291    }
292
293    /// Sets the returning clause on the query's `SELECT` body.
294    pub fn returning(mut self, returning: Returning) -> Self {
295        match &mut self.query.body {
296            ExprSet::Select(select) => {
297                select.returning = returning;
298            }
299            _ => todo!(),
300        }
301
302        self
303    }
304
305    /// Sets the returning clause to `Returning::Project` containing the given
306    /// expression.
307    pub fn returning_project(self, expr: impl Into<Expr>) -> Self {
308        self.returning(Returning::Project(expr.into()))
309    }
310
311    /// Sets the returning clause to `Returning::Expr` containing the given
312    /// expression.
313    pub fn returning_expr(self, expr: impl Into<Expr>) -> Self {
314        self.returning(Returning::Expr(expr.into()))
315    }
316
317    /// Consumes this builder and returns the constructed [`Query`].
318    pub fn build(self) -> Query {
319        self.query
320    }
321}
322
323impl From<QueryBuilder> for Query {
324    fn from(value: QueryBuilder) -> Self {
325        value.build()
326    }
327}