Skip to main content

dbkit_core/
query.rs

1use std::marker::PhantomData;
2
3use crate::compile::{CompiledSql, SqlBuilder, ToSql};
4use crate::expr::{Expr, ExprNode, IntoExpr};
5use crate::func;
6use crate::load::{ApplyLoad, LoadChain, NoLoad};
7use crate::rel::RelationInfo;
8use crate::schema::{ColumnRef, Table};
9
10#[derive(Debug, Clone, Copy)]
11pub enum JoinKind {
12    Inner,
13    Left,
14}
15
16#[derive(Debug, Clone)]
17pub struct Join {
18    pub table: Table,
19    pub on: Expr<bool>,
20    pub kind: JoinKind,
21}
22
23#[derive(Debug, Clone)]
24pub struct SelectItem {
25    pub expr: ExprNode,
26    pub alias: Option<String>,
27}
28
29#[derive(Debug, Clone, Copy)]
30pub enum OrderDirection {
31    Asc,
32    Desc,
33}
34
35#[derive(Debug, Clone)]
36pub enum OrderExpr {
37    Expr(ExprNode),
38    Alias(String),
39}
40
41pub trait IntoOrderExpr {
42    fn into_order_expr(self) -> OrderExpr;
43}
44
45impl IntoOrderExpr for ColumnRef {
46    fn into_order_expr(self) -> OrderExpr {
47        OrderExpr::Expr(ExprNode::Column(self))
48    }
49}
50
51impl<M, T> IntoOrderExpr for crate::schema::Column<M, T> {
52    fn into_order_expr(self) -> OrderExpr {
53        OrderExpr::Expr(ExprNode::Column(self.as_ref()))
54    }
55}
56
57impl<T, Kind> IntoOrderExpr for Expr<T, Kind> {
58    fn into_order_expr(self) -> OrderExpr {
59        OrderExpr::Expr(self.node)
60    }
61}
62
63#[derive(Debug, Clone)]
64pub struct Order {
65    pub expr: OrderExpr,
66    pub direction: OrderDirection,
67}
68
69#[derive(Debug, Clone, Copy, Default)]
70pub struct NoRowLock;
71
72#[derive(Debug, Clone, Copy, Default)]
73pub struct ForUpdateRowLock;
74
75#[derive(Debug, Clone, Copy, Default)]
76pub struct NotDistinct;
77
78#[derive(Debug, Clone, Copy, Default)]
79pub struct DistinctSelected;
80
81#[derive(Debug, Clone, Copy, Default)]
82pub struct NotGrouped;
83
84#[derive(Debug, Clone, Copy, Default)]
85pub struct Grouped;
86
87#[derive(Debug, Clone, Copy, PartialEq, Eq)]
88enum RowLockWait {
89    Wait,
90    SkipLocked,
91    NoWait,
92}
93
94#[derive(Debug, Clone)]
95pub struct Select<Out, Loads = NoLoad, Lock = NoRowLock, DistinctState = NotDistinct, GroupState = NotGrouped> {
96    table: Table,
97    columns: Option<Vec<SelectItem>>,
98    joins: Vec<Join>,
99    filters: Vec<Expr<bool>>,
100    group_by: Vec<ExprNode>,
101    having: Vec<Expr<bool>>,
102    order_by: Vec<Order>,
103    limit: Option<u64>,
104    offset: Option<u64>,
105    distinct: bool,
106    row_lock_wait: Option<RowLockWait>,
107    loads: Loads,
108    _marker: PhantomData<Out>,
109    _lock_marker: PhantomData<Lock>,
110    _distinct_marker: PhantomData<DistinctState>,
111    _group_marker: PhantomData<GroupState>,
112}
113
114impl<Out> Select<Out, NoLoad, NoRowLock, NotDistinct, NotGrouped> {
115    pub fn new(table: Table) -> Self {
116        Self {
117            table,
118            columns: None,
119            joins: Vec::new(),
120            filters: Vec::new(),
121            group_by: Vec::new(),
122            having: Vec::new(),
123            order_by: Vec::new(),
124            limit: None,
125            offset: None,
126            distinct: false,
127            row_lock_wait: None,
128            loads: NoLoad,
129            _marker: PhantomData,
130            _lock_marker: PhantomData,
131            _distinct_marker: PhantomData,
132            _group_marker: PhantomData,
133        }
134    }
135}
136
137impl<Out, Loads, Lock, DistinctState, GroupState> Select<Out, Loads, Lock, DistinctState, GroupState> {
138    pub fn table(&self) -> Table {
139        self.table
140    }
141
142    pub fn columns_ref(&self) -> Option<&[SelectItem]> {
143        self.columns.as_deref()
144    }
145
146    pub fn joins(&self) -> &[Join] {
147        &self.joins
148    }
149
150    pub fn select_only(mut self) -> Self {
151        self.columns = Some(Vec::new());
152        self
153    }
154
155    pub fn column<T>(mut self, expr: impl IntoExpr<T>) -> Self {
156        let item = SelectItem {
157            expr: expr.into_expr().node,
158            alias: None,
159        };
160        match &mut self.columns {
161            Some(columns) => columns.push(item),
162            None => self.columns = Some(vec![item]),
163        }
164        self
165    }
166
167    pub fn column_as<T>(mut self, expr: impl IntoExpr<T>, alias: &str) -> Self {
168        let item = SelectItem {
169            expr: expr.into_expr().node,
170            alias: Some(alias.to_string()),
171        };
172        match &mut self.columns {
173            Some(columns) => columns.push(item),
174            None => self.columns = Some(vec![item]),
175        }
176        self
177    }
178
179    pub fn filter(mut self, expr: Expr<bool>) -> Self {
180        self.filters.push(expr);
181        self
182    }
183
184    pub fn where_exists<SubOut, SubLoads, SubLock, SubDistinctState, SubGroupState>(
185        self,
186        subquery: Select<SubOut, SubLoads, SubLock, SubDistinctState, SubGroupState>,
187    ) -> Self {
188        self.filter(func::exists(subquery))
189    }
190
191    pub fn where_not_exists<SubOut, SubLoads, SubLock, SubDistinctState, SubGroupState>(
192        self,
193        subquery: Select<SubOut, SubLoads, SubLock, SubDistinctState, SubGroupState>,
194    ) -> Self {
195        self.filter(func::exists(subquery).not())
196    }
197
198    pub fn join<R>(mut self, rel: R) -> Self
199    where
200        R: RelationInfo<Parent = Out>,
201    {
202        let relation = rel.relation();
203        for (table, on) in relation.join_steps() {
204            self.joins.push(Join {
205                table,
206                on,
207                kind: JoinKind::Inner,
208            });
209        }
210        self
211    }
212
213    pub fn left_join<R>(mut self, rel: R) -> Self
214    where
215        R: RelationInfo<Parent = Out>,
216    {
217        let relation = rel.relation();
218        for (table, on) in relation.join_steps() {
219            self.joins.push(Join {
220                table,
221                on,
222                kind: JoinKind::Left,
223            });
224        }
225        self
226    }
227
228    pub fn join_on(mut self, table: Table, on: Expr<bool>) -> Self {
229        self.joins.push(Join {
230            table,
231            on,
232            kind: JoinKind::Inner,
233        });
234        self
235    }
236
237    pub fn left_join_on(mut self, table: Table, on: Expr<bool>) -> Self {
238        self.joins.push(Join {
239            table,
240            on,
241            kind: JoinKind::Left,
242        });
243        self
244    }
245
246    pub fn limit(mut self, limit: u64) -> Self {
247        self.limit = Some(limit);
248        self
249    }
250
251    pub fn offset(mut self, offset: u64) -> Self {
252        self.offset = Some(offset);
253        self
254    }
255
256    pub fn order_by(mut self, order: Order) -> Self {
257        self.order_by.push(order);
258        self
259    }
260
261    pub fn columns(mut self, columns: Vec<ColumnRef>) -> Self {
262        let items = columns
263            .into_iter()
264            .map(|col| SelectItem {
265                expr: ExprNode::Column(col),
266                alias: None,
267            })
268            .collect::<Vec<_>>();
269        self.columns = Some(items);
270        self
271    }
272
273    pub fn into_model<T>(self) -> Select<T, Loads, Lock, DistinctState, GroupState> {
274        Select {
275            table: self.table,
276            columns: self.columns,
277            joins: self.joins,
278            filters: self.filters,
279            group_by: self.group_by,
280            having: self.having,
281            order_by: self.order_by,
282            limit: self.limit,
283            offset: self.offset,
284            distinct: self.distinct,
285            row_lock_wait: self.row_lock_wait,
286            loads: self.loads,
287            _marker: PhantomData,
288            _lock_marker: PhantomData,
289            _distinct_marker: PhantomData,
290            _group_marker: PhantomData,
291        }
292    }
293
294    pub fn with<L>(self, load: L) -> Select<L::Out2, LoadChain<Loads, L>, Lock, DistinctState, GroupState>
295    where
296        L: ApplyLoad<Out>,
297    {
298        Select {
299            table: self.table,
300            columns: self.columns,
301            joins: self.joins,
302            filters: self.filters,
303            group_by: self.group_by,
304            having: self.having,
305            order_by: self.order_by,
306            limit: self.limit,
307            offset: self.offset,
308            distinct: self.distinct,
309            row_lock_wait: self.row_lock_wait,
310            loads: LoadChain { prev: self.loads, load },
311            _marker: PhantomData,
312            _lock_marker: PhantomData,
313            _distinct_marker: PhantomData,
314            _group_marker: PhantomData,
315        }
316    }
317
318    pub fn compile(&self) -> CompiledSql {
319        self.compile_inner(true, true, true)
320    }
321
322    pub fn compile_without_pagination(&self) -> CompiledSql {
323        self.compile_inner(false, false, false)
324    }
325
326    pub(crate) fn compile_for_exists(&self) -> CompiledSql {
327        self.compile_inner(true, true, true)
328    }
329
330    pub fn compile_with_extra(&self, extra_columns: &[SelectItem], extra_joins: &[Join]) -> CompiledSql {
331        self.compile_inner_with(extra_columns, extra_joins, true, true, true)
332    }
333
334    fn compile_inner(&self, include_order: bool, include_pagination: bool, include_locking: bool) -> CompiledSql {
335        self.compile_inner_with(&[], &[], include_order, include_pagination, include_locking)
336    }
337
338    fn compile_inner_with(
339        &self,
340        extra_columns: &[SelectItem],
341        extra_joins: &[Join],
342        include_order: bool,
343        include_pagination: bool,
344        include_locking: bool,
345    ) -> CompiledSql {
346        let mut builder = SqlBuilder::new();
347        builder.push_sql("SELECT ");
348        if self.distinct {
349            builder.push_sql("DISTINCT ");
350        }
351        match &self.columns {
352            Some(columns) => {
353                for (idx, col) in columns.iter().enumerate() {
354                    if idx > 0 {
355                        builder.push_sql(", ");
356                    }
357                    col.expr.to_sql(&mut builder);
358                    if let Some(alias) = &col.alias {
359                        builder.push_sql(" AS ");
360                        builder.push_sql(alias);
361                    }
362                }
363                if !extra_columns.is_empty() {
364                    for col in extra_columns {
365                        builder.push_sql(", ");
366                        col.expr.to_sql(&mut builder);
367                        if let Some(alias) = &col.alias {
368                            builder.push_sql(" AS ");
369                            builder.push_sql(alias);
370                        }
371                    }
372                }
373            }
374            None => {
375                builder.push_sql(self.table.qualifier());
376                builder.push_sql(".*");
377                if !extra_columns.is_empty() {
378                    for col in extra_columns {
379                        builder.push_sql(", ");
380                        col.expr.to_sql(&mut builder);
381                        if let Some(alias) = &col.alias {
382                            builder.push_sql(" AS ");
383                            builder.push_sql(alias);
384                        }
385                    }
386                }
387            }
388        }
389        builder.push_sql(" FROM ");
390        builder.push_sql(&self.table.qualified_name());
391        if let Some(alias) = self.table.alias {
392            builder.push_sql(" ");
393            builder.push_sql(alias);
394        }
395        for join in &self.joins {
396            builder.push_sql(match join.kind {
397                JoinKind::Inner => " JOIN ",
398                JoinKind::Left => " LEFT JOIN ",
399            });
400            builder.push_sql(&join.table.qualified_name());
401            if let Some(alias) = join.table.alias {
402                builder.push_sql(" ");
403                builder.push_sql(alias);
404            }
405            builder.push_sql(" ON ");
406            join.on.node.to_sql(&mut builder);
407        }
408        for join in extra_joins {
409            builder.push_sql(match join.kind {
410                JoinKind::Inner => " JOIN ",
411                JoinKind::Left => " LEFT JOIN ",
412            });
413            builder.push_sql(&join.table.qualified_name());
414            if let Some(alias) = join.table.alias {
415                builder.push_sql(" ");
416                builder.push_sql(alias);
417            }
418            builder.push_sql(" ON ");
419            join.on.node.to_sql(&mut builder);
420        }
421        if !self.filters.is_empty() {
422            builder.push_sql(" WHERE ");
423            for (idx, expr) in self.filters.iter().enumerate() {
424                if idx > 0 {
425                    builder.push_sql(" AND ");
426                }
427                expr.node.to_sql(&mut builder);
428            }
429        }
430        if !self.group_by.is_empty() {
431            builder.push_sql(" GROUP BY ");
432            for (idx, expr) in self.group_by.iter().enumerate() {
433                if idx > 0 {
434                    builder.push_sql(", ");
435                }
436                expr.to_sql(&mut builder);
437            }
438        }
439        if !self.having.is_empty() {
440            builder.push_sql(" HAVING ");
441            for (idx, expr) in self.having.iter().enumerate() {
442                if idx > 0 {
443                    builder.push_sql(" AND ");
444                }
445                expr.node.to_sql(&mut builder);
446            }
447        }
448        if include_order && !self.order_by.is_empty() {
449            builder.push_sql(" ORDER BY ");
450            for (idx, order) in self.order_by.iter().enumerate() {
451                if idx > 0 {
452                    builder.push_sql(", ");
453                }
454                match &order.expr {
455                    OrderExpr::Expr(expr) => expr.to_sql(&mut builder),
456                    OrderExpr::Alias(alias) => builder.push_sql(alias),
457                }
458                builder.push_sql(match order.direction {
459                    OrderDirection::Asc => " ASC",
460                    OrderDirection::Desc => " DESC",
461                });
462            }
463        }
464        if include_pagination {
465            if let Some(limit) = self.limit {
466                builder.push_sql(" LIMIT ");
467                builder.push_sql(&limit.to_string());
468            }
469            if let Some(offset) = self.offset {
470                builder.push_sql(" OFFSET ");
471                builder.push_sql(&offset.to_string());
472            }
473        }
474        if include_locking {
475            if let Some(wait) = self.row_lock_wait {
476                builder.push_sql(" FOR UPDATE");
477                if self
478                    .joins
479                    .iter()
480                    .chain(extra_joins.iter())
481                    .any(|join| matches!(join.kind, JoinKind::Left))
482                {
483                    builder.push_sql(" OF ");
484                    builder.push_sql(self.table.qualifier());
485                }
486                match wait {
487                    RowLockWait::Wait => {}
488                    RowLockWait::SkipLocked => builder.push_sql(" SKIP LOCKED"),
489                    RowLockWait::NoWait => builder.push_sql(" NOWAIT"),
490                }
491            }
492        }
493        builder.finish()
494    }
495
496    pub fn debug_sql(&self) -> String {
497        self.compile().sql
498    }
499
500    pub fn into_parts(self) -> (CompiledSql, Loads) {
501        let compiled = self.compile();
502        (compiled, self.loads)
503    }
504
505    pub fn into_parts_with_loads(self) -> (Select<Out, NoLoad, Lock, DistinctState, GroupState>, Loads) {
506        let Select {
507            table,
508            columns,
509            joins,
510            filters,
511            group_by,
512            having,
513            order_by,
514            limit,
515            offset,
516            distinct,
517            row_lock_wait,
518            loads,
519            _marker,
520            _lock_marker,
521            _distinct_marker,
522            _group_marker,
523        } = self;
524
525        let select = Select {
526            table,
527            columns,
528            joins,
529            filters,
530            group_by,
531            having,
532            order_by,
533            limit,
534            offset,
535            distinct,
536            row_lock_wait,
537            loads: NoLoad,
538            _marker,
539            _lock_marker: PhantomData,
540            _distinct_marker: PhantomData,
541            _group_marker: PhantomData,
542        };
543
544        (select, loads)
545    }
546}
547
548impl<Out, Loads, DistinctState, GroupState> Select<Out, Loads, NoRowLock, DistinctState, GroupState> {
549    pub fn group_by<T>(mut self, expr: impl IntoExpr<T>) -> Select<Out, Loads, NoRowLock, DistinctState, Grouped> {
550        self.group_by.push(expr.into_expr().node);
551        Select {
552            table: self.table,
553            columns: self.columns,
554            joins: self.joins,
555            filters: self.filters,
556            group_by: self.group_by,
557            having: self.having,
558            order_by: self.order_by,
559            limit: self.limit,
560            offset: self.offset,
561            distinct: self.distinct,
562            row_lock_wait: self.row_lock_wait,
563            loads: self.loads,
564            _marker: PhantomData,
565            _lock_marker: PhantomData,
566            _distinct_marker: PhantomData,
567            _group_marker: PhantomData,
568        }
569    }
570
571    pub fn having(mut self, expr: Expr<bool>) -> Select<Out, Loads, NoRowLock, DistinctState, Grouped> {
572        self.having.push(expr);
573        Select {
574            table: self.table,
575            columns: self.columns,
576            joins: self.joins,
577            filters: self.filters,
578            group_by: self.group_by,
579            having: self.having,
580            order_by: self.order_by,
581            limit: self.limit,
582            offset: self.offset,
583            distinct: self.distinct,
584            row_lock_wait: self.row_lock_wait,
585            loads: self.loads,
586            _marker: PhantomData,
587            _lock_marker: PhantomData,
588            _distinct_marker: PhantomData,
589            _group_marker: PhantomData,
590        }
591    }
592}
593
594impl<Out, Loads, GroupState> Select<Out, Loads, NoRowLock, NotDistinct, GroupState> {
595    pub fn distinct(mut self) -> Select<Out, Loads, NoRowLock, DistinctSelected, GroupState> {
596        self.distinct = true;
597        Select {
598            table: self.table,
599            columns: self.columns,
600            joins: self.joins,
601            filters: self.filters,
602            group_by: self.group_by,
603            having: self.having,
604            order_by: self.order_by,
605            limit: self.limit,
606            offset: self.offset,
607            distinct: self.distinct,
608            row_lock_wait: self.row_lock_wait,
609            loads: self.loads,
610            _marker: PhantomData,
611            _lock_marker: PhantomData,
612            _distinct_marker: PhantomData,
613            _group_marker: PhantomData,
614        }
615    }
616}
617
618impl<Out, Loads> Select<Out, Loads, NoRowLock, NotDistinct, NotGrouped> {
619    pub fn for_update(self) -> Select<Out, Loads, ForUpdateRowLock, NotDistinct, NotGrouped> {
620        Select {
621            table: self.table,
622            columns: self.columns,
623            joins: self.joins,
624            filters: self.filters,
625            group_by: self.group_by,
626            having: self.having,
627            order_by: self.order_by,
628            limit: self.limit,
629            offset: self.offset,
630            distinct: self.distinct,
631            row_lock_wait: Some(self.row_lock_wait.unwrap_or(RowLockWait::Wait)),
632            loads: self.loads,
633            _marker: PhantomData,
634            _lock_marker: PhantomData,
635            _distinct_marker: PhantomData,
636            _group_marker: PhantomData,
637        }
638    }
639}
640
641impl<Out, Loads, GroupState> Select<Out, Loads, NoRowLock, DistinctSelected, GroupState> {
642    pub fn distinct(self) -> Self {
643        self
644    }
645}
646
647impl<Out, Loads> Select<Out, Loads, ForUpdateRowLock, NotDistinct, NotGrouped> {
648    pub fn for_update(self) -> Self {
649        self
650    }
651
652    pub fn skip_locked(mut self) -> Self {
653        self.row_lock_wait = Some(RowLockWait::SkipLocked);
654        self
655    }
656
657    pub fn nowait(mut self) -> Self {
658        self.row_lock_wait = Some(RowLockWait::NoWait);
659        self
660    }
661}
662
663impl Order {
664    pub fn asc(expr: impl IntoOrderExpr) -> Self {
665        Self {
666            expr: expr.into_order_expr(),
667            direction: OrderDirection::Asc,
668        }
669    }
670
671    pub fn desc(expr: impl IntoOrderExpr) -> Self {
672        Self {
673            expr: expr.into_order_expr(),
674            direction: OrderDirection::Desc,
675        }
676    }
677
678    pub fn asc_alias(alias: &str) -> Self {
679        Self {
680            expr: OrderExpr::Alias(alias.to_string()),
681            direction: OrderDirection::Asc,
682        }
683    }
684
685    pub fn desc_alias(alias: &str) -> Self {
686        Self {
687            expr: OrderExpr::Alias(alias.to_string()),
688            direction: OrderDirection::Desc,
689        }
690    }
691}