Skip to main content

a3s_orm/query/
select.rs

1use std::marker::PhantomData;
2
3use crate::ast::{
4    JoinKind, JoinNode, QueryNode, SelectLockNode, SelectLockStrength, SelectLockWait, SelectNode,
5    SetOperationKind, SetOperationNode, TableNode,
6};
7use crate::expression::{Column, Expression, OrderDirection, Selection};
8use crate::function::TypedExpression;
9use crate::schema::{Table, TableRef};
10
11use super::{Cte, Query};
12
13#[derive(Clone, Debug)]
14pub struct SelectQuery<T: Table, O = ()> {
15    node: SelectNode,
16    marker: PhantomData<fn() -> (T, O)>,
17}
18
19pub fn select_from<T: Table>() -> SelectQuery<T> {
20    SelectQuery::new(TableRef::<T>::new())
21}
22
23pub fn select_from_as<Source: Table, Alias: Table>() -> SelectQuery<Alias> {
24    SelectQuery::from_table(TableNode {
25        name: Source::NAME,
26        alias: Some(Alias::NAME),
27    })
28}
29
30impl<T: Table> SelectQuery<T> {
31    pub(crate) fn new(table: TableRef<T>) -> Self {
32        Self::from_table(table_node(table))
33    }
34
35    fn from_table(from: TableNode) -> Self {
36        Self {
37            node: SelectNode {
38                ctes: Vec::new(),
39                from,
40                selections: Vec::new(),
41                joins: Vec::new(),
42                filter: None,
43                group_by: Vec::new(),
44                having: None,
45                set_operations: Vec::new(),
46                order_by: Vec::new(),
47                limit: None,
48                offset: None,
49                distinct: false,
50                lock: None,
51            },
52            marker: PhantomData,
53        }
54    }
55}
56
57impl<T: Table, O> SelectQuery<T, O> {
58    pub fn select<S: Selection>(mut self, selection: S) -> SelectQuery<T, S::Output> {
59        self.node.selections = selection.expressions();
60        SelectQuery {
61            node: self.node,
62            marker: PhantomData,
63        }
64    }
65
66    /// Select every column for driver-owned row access.
67    ///
68    /// The output is intentionally untyped because a table marker does not
69    /// define a Rust record decoder. Use an explicit selection with
70    /// `fetch_all_as` when typed decoding is required.
71    pub fn select_all(mut self) -> SelectQuery<T, ()> {
72        self.node.selections = vec![Expression::Column {
73            table: T::NAME,
74            name: "*",
75        }];
76        SelectQuery {
77            node: self.node,
78            marker: PhantomData,
79        }
80    }
81
82    pub fn distinct(mut self) -> Self {
83        self.node.distinct = true;
84        self
85    }
86
87    pub fn with<C: Table>(mut self, cte: Cte<C>) -> Self {
88        self.node.ctes.push(cte.node);
89        self
90    }
91
92    pub fn as_cte<C: Table>(self) -> Cte<C> {
93        Cte::new(self.node)
94    }
95
96    pub fn filter(mut self, expression: Expression) -> Self {
97        self.node.filter = Some(match self.node.filter.take() {
98            Some(existing) => existing.and(expression),
99            None => expression,
100        });
101        self
102    }
103
104    pub fn group_by<TableType, ValueType>(mut self, column: Column<TableType, ValueType>) -> Self {
105        self.node.group_by.push(column.expression());
106        self
107    }
108
109    pub fn having(mut self, expression: Expression) -> Self {
110        self.node.having = Some(match self.node.having.take() {
111            Some(existing) => existing.and(expression),
112            None => expression,
113        });
114        self
115    }
116
117    pub fn inner_join<J: Table>(self, on: Expression) -> Self {
118        self.join::<J>(JoinKind::Inner, on)
119    }
120
121    pub fn left_join<J: Table>(self, on: Expression) -> Self {
122        self.join::<J>(JoinKind::Left, on)
123    }
124
125    pub fn right_join<J: Table>(self, on: Expression) -> Self {
126        self.join::<J>(JoinKind::Right, on)
127    }
128
129    pub fn full_join<J: Table>(self, on: Expression) -> Self {
130        self.join::<J>(JoinKind::Full, on)
131    }
132
133    pub fn inner_join_as<Source: Table, Alias: Table>(self, on: Expression) -> Self {
134        self.join_as::<Source, Alias>(JoinKind::Inner, on)
135    }
136
137    pub fn left_join_as<Source: Table, Alias: Table>(self, on: Expression) -> Self {
138        self.join_as::<Source, Alias>(JoinKind::Left, on)
139    }
140
141    pub fn order_by<TableType, ValueType>(
142        mut self,
143        column: Column<TableType, ValueType>,
144        direction: OrderDirection,
145    ) -> Self {
146        self.node.order_by.push((column.expression(), direction));
147        self
148    }
149
150    pub fn order_by_expression<ValueType>(
151        mut self,
152        expression: TypedExpression<ValueType>,
153        direction: OrderDirection,
154    ) -> Self {
155        self.node
156            .order_by
157            .push((expression.expression(), direction));
158        self
159    }
160
161    pub fn limit(mut self, limit: u64) -> Self {
162        self.node.limit = Some(limit);
163        self
164    }
165
166    pub fn offset(mut self, offset: u64) -> Self {
167        self.node.offset = Some(offset);
168        self
169    }
170
171    /// Lock every selected row with PostgreSQL `FOR UPDATE`.
172    pub fn for_update(self) -> Self {
173        self.lock_rows(SelectLockStrength::Update)
174    }
175
176    /// Lock only the named source or join marker with PostgreSQL
177    /// `FOR UPDATE OF`.
178    pub fn for_update_of<Locked: Table>(self) -> Self {
179        self.lock_rows_of::<Locked>(SelectLockStrength::Update)
180    }
181
182    /// Lock every selected row with PostgreSQL `FOR NO KEY UPDATE`.
183    pub fn for_no_key_update(self) -> Self {
184        self.lock_rows(SelectLockStrength::NoKeyUpdate)
185    }
186
187    /// Lock only the named source or join marker with PostgreSQL
188    /// `FOR NO KEY UPDATE OF`.
189    pub fn for_no_key_update_of<Locked: Table>(self) -> Self {
190        self.lock_rows_of::<Locked>(SelectLockStrength::NoKeyUpdate)
191    }
192
193    /// Lock every selected row with PostgreSQL `FOR SHARE`.
194    pub fn for_share(self) -> Self {
195        self.lock_rows(SelectLockStrength::Share)
196    }
197
198    /// Lock only the named source or join marker with PostgreSQL `FOR SHARE OF`.
199    pub fn for_share_of<Locked: Table>(self) -> Self {
200        self.lock_rows_of::<Locked>(SelectLockStrength::Share)
201    }
202
203    /// Lock every selected row with PostgreSQL `FOR KEY SHARE`.
204    pub fn for_key_share(self) -> Self {
205        self.lock_rows(SelectLockStrength::KeyShare)
206    }
207
208    /// Lock only the named source or join marker with PostgreSQL
209    /// `FOR KEY SHARE OF`.
210    pub fn for_key_share_of<Locked: Table>(self) -> Self {
211        self.lock_rows_of::<Locked>(SelectLockStrength::KeyShare)
212    }
213
214    /// Fail instead of waiting for a conflicting row lock.
215    pub fn no_wait(mut self) -> Self {
216        self.lock_mut().wait = SelectLockWait::NoWait;
217        self
218    }
219
220    /// Skip rows currently held by a conflicting row lock.
221    pub fn skip_locked(mut self) -> Self {
222        self.lock_mut().wait = SelectLockWait::SkipLocked;
223        self
224    }
225
226    pub fn union<Source: Table>(self, query: SelectQuery<Source, O>) -> Self {
227        self.set_operation(SetOperationKind::Union, query)
228    }
229
230    pub fn union_all<Source: Table>(self, query: SelectQuery<Source, O>) -> Self {
231        self.set_operation(SetOperationKind::UnionAll, query)
232    }
233
234    pub fn intersect<Source: Table>(self, query: SelectQuery<Source, O>) -> Self {
235        self.set_operation(SetOperationKind::Intersect, query)
236    }
237
238    pub fn except<Source: Table>(self, query: SelectQuery<Source, O>) -> Self {
239        self.set_operation(SetOperationKind::Except, query)
240    }
241
242    fn join<J: Table>(mut self, kind: JoinKind, on: Expression) -> Self {
243        self.node.joins.push(JoinNode {
244            kind,
245            table: table_node(TableRef::<J>::new()),
246            on,
247        });
248        self
249    }
250
251    fn join_as<Source: Table, Alias: Table>(mut self, kind: JoinKind, on: Expression) -> Self {
252        self.node.joins.push(JoinNode {
253            kind,
254            table: TableNode {
255                name: Source::NAME,
256                alias: Some(Alias::NAME),
257            },
258            on,
259        });
260        self
261    }
262
263    fn set_operation<Source: Table>(
264        mut self,
265        kind: SetOperationKind,
266        query: SelectQuery<Source, O>,
267    ) -> Self {
268        self.node.set_operations.push(SetOperationNode {
269            kind,
270            query: Box::new(query.node),
271        });
272        self
273    }
274
275    fn lock_mut(&mut self) -> &mut SelectLockNode {
276        self.node.lock.get_or_insert_with(|| SelectLockNode {
277            strength: SelectLockStrength::Update,
278            tables: Vec::new(),
279            wait: SelectLockWait::Block,
280        })
281    }
282
283    fn lock_rows(mut self, strength: SelectLockStrength) -> Self {
284        self.node.lock = Some(SelectLockNode {
285            strength,
286            tables: Vec::new(),
287            wait: SelectLockWait::Block,
288        });
289        self
290    }
291
292    fn lock_rows_of<Locked: Table>(mut self, strength: SelectLockStrength) -> Self {
293        let lock = self.lock_mut();
294        if lock.strength != strength {
295            lock.strength = strength;
296            lock.tables.clear();
297            lock.wait = SelectLockWait::Block;
298        }
299        if !lock.tables.contains(&Locked::NAME) {
300            lock.tables.push(Locked::NAME);
301        }
302        self
303    }
304
305    pub(crate) fn into_node(self) -> SelectNode {
306        self.node
307    }
308}
309
310impl<T: Table, O> Query for SelectQuery<T, O> {
311    type Output = O;
312
313    fn compile(self, dialect: &impl crate::Dialect) -> crate::Result<crate::CompiledQuery> {
314        crate::compiler::compile(QueryNode::Select(Box::new(self.node)), dialect)
315    }
316}
317
318fn table_node<T: Table>(table: TableRef<T>) -> TableNode {
319    TableNode {
320        name: table.name(),
321        alias: None,
322    }
323}