Skip to main content

qbrs_core/select/
set_op.rs

1//! `UNION`/`UNION ALL`/`INTERSECT`/`EXCEPT` between two `SELECT`s that may
2//! have entirely different `Scope`s (different tables, different JOINs) —
3//! the only thing that must line up is their *output shape*: `row::SameShape`
4//! requires the same column names, in the same order, decoding to the same
5//! types. Names as well as types, because the combined result is read by key
6//! — a branch whose columns merely happen to be type-compatible would
7//! otherwise splice in transposed. Keys from different tables still match,
8//! since the comparison is on names, and SQL itself takes a `UNION`'s column
9//! names from the first branch.
10
11use std::marker::PhantomData;
12
13use super::{Select, Selection, SortDir};
14use crate::dialect::Dialect;
15use crate::expr::Value;
16use crate::render::{Fragment, QuerySink, Sink};
17use crate::row::SameShape;
18
19#[derive(Clone, Copy, PartialEq, Eq)]
20enum SetOpKind {
21    Union,
22    UnionAll,
23    Intersect,
24    Except,
25}
26
27impl SetOpKind {
28    fn keyword(&self) -> &'static str {
29        match self {
30            SetOpKind::Union => " UNION ",
31            SetOpKind::UnionAll => " UNION ALL ",
32            SetOpKind::Intersect => " INTERSECT ",
33            SetOpKind::Except => " EXCEPT ",
34        }
35    }
36}
37
38/// A chain of `SELECT`s combined by set operators, all decoding to the first
39/// branch's `Output` — which is also where SQL itself takes the combined
40/// result's column names from. `ORDER BY` here is necessarily by **ordinal position**
41/// (`ORDER BY 1`, 1-indexed) rather than a typed column — the branches can
42/// have entirely different `Scope`s, so there is no single scope left to
43/// check a column reference against once they're combined; ordinal position
44/// is the only reference SQL itself allows in this position.
45pub struct SetOp<D, Output> {
46    first: Fragment,
47    rest: Vec<(SetOpKind, Fragment)>,
48    order_by: Vec<(u32, SortDir)>,
49    limit: Option<super::RowCount>,
50    offset: Option<super::RowCount>,
51    _marker: PhantomData<fn() -> (D, Output)>,
52}
53
54impl<D: Dialect, L> SetOp<D, crate::row::Row<L>> {
55    /// `ORDER BY` naming the column instead of counting to it: the position
56    /// is `row::Field`'s index, which the row already carries. A column the
57    /// combined result doesn't select is a compile error, which is the whole
58    /// reason no `ORDER BY <n>` is spellable here. Callable multiple times
59    /// like `Select::order_by`, each call appending a key.
60    pub fn order_by_column<K, Idx>(self, _key: K, dir: SortDir) -> Self
61    where
62        K: crate::row::LookupKey,
63        L: crate::row::Field<K::Key, Idx>,
64        Idx: crate::scope::Position,
65    {
66        self.order_by_ordinal(<Idx as crate::scope::Position>::POSITION, dir)
67    }
68}
69
70/// A set operation whose branches select one un-tupled column: its output
71/// is a bare value, so there is one position and nothing to name.
72impl<D: Dialect, V: crate::select::SingleColumn> SetOp<D, V> {
73    pub fn order_by(self, dir: SortDir) -> Self {
74        self.order_by_ordinal(1, dir)
75    }
76}
77
78impl<D: Dialect, Output> SetOp<D, Output> {
79    fn new(first: Fragment) -> Self {
80        SetOp {
81            first,
82            rest: Vec::new(),
83            order_by: Vec::new(),
84            limit: None,
85            offset: None,
86            _marker: PhantomData,
87        }
88    }
89
90    fn push(mut self, kind: SetOpKind, branch: Fragment) -> Self {
91        self.rest.push((kind, branch));
92        self
93    }
94
95    /// Appends another branch via `UNION` (duplicates across branches are
96    /// removed, same as plain SQL `UNION`).
97    pub fn union<ScopeB, SelB, IdxB>(self, other: &Select<D, ScopeB, SelB>) -> Self
98    where
99        SelB: Selection<ScopeB, IdxB>,
100        SelB::Output: SameShape<Output>,
101    {
102        self.push(SetOpKind::Union, other.fragment::<IdxB>())
103    }
104
105    /// Appends another branch via `UNION ALL` (no deduplication — cheaper
106    /// than `UNION` when the branches are already known disjoint, or when
107    /// duplicates are meaningful).
108    pub fn union_all<ScopeB, SelB, IdxB>(self, other: &Select<D, ScopeB, SelB>) -> Self
109    where
110        SelB: Selection<ScopeB, IdxB>,
111        SelB::Output: SameShape<Output>,
112    {
113        self.push(SetOpKind::UnionAll, other.fragment::<IdxB>())
114    }
115
116    /// Appends another branch via `INTERSECT` (rows present in both).
117    pub fn intersect<ScopeB, SelB, IdxB>(self, other: &Select<D, ScopeB, SelB>) -> Self
118    where
119        SelB: Selection<ScopeB, IdxB>,
120        SelB::Output: SameShape<Output>,
121    {
122        self.push(SetOpKind::Intersect, other.fragment::<IdxB>())
123    }
124
125    /// Appends another branch via `EXCEPT` (rows in the accumulated result
126    /// so far, minus rows in `other`).
127    pub fn except<ScopeB, SelB, IdxB>(self, other: &Select<D, ScopeB, SelB>) -> Self
128    where
129        SelB: Selection<ScopeB, IdxB>,
130        SelB::Output: SameShape<Output>,
131    {
132        self.push(SetOpKind::Except, other.fragment::<IdxB>())
133    }
134
135    fn order_by_ordinal(mut self, position: u32, dir: SortDir) -> Self {
136        self.order_by.push((position, dir));
137        self
138    }
139
140    pub fn limit(mut self, n: impl super::IntoRowCount) -> Self {
141        self.limit = Some(n.into_row_count());
142        self
143    }
144
145    pub fn offset(mut self, n: impl super::IntoRowCount) -> Self {
146        self.offset = Some(n.into_row_count());
147        self
148    }
149
150    /// How many rows the combination returns, its own `ORDER BY`/paging
151    /// dropped. The branches keep theirs: a `UNION` of two `LIMIT`ed queries
152    /// is a different set from a `UNION` of the whole ones.
153    pub fn count_sql(&self, _dialect: D) -> (String, Vec<Value>) {
154        let mut sink = QuerySink::<D>::new();
155        crate::render::render_count_wrapped::<D>(&mut sink, |sink| self.render_branches(sink));
156        sink.finish()
157    }
158
159    pub fn to_sql(&self, _dialect: D) -> (String, Vec<Value>) {
160        let mut sink = QuerySink::<D>::new();
161        self.render_branches(&mut sink);
162        self.render_ordering(&mut sink);
163        sink.finish()
164    }
165
166    /// The set operation itself, without the ordering and paging applied to
167    /// its result — which is what a count of it must leave out.
168    fn render_branches(&self, sink: &mut QuerySink<D>) {
169        let branch = |sink: &mut QuerySink<D>, sql: &Fragment| {
170            // A branch has to be shut off from the operator beside it:
171            // otherwise its `ORDER BY`/`LIMIT`, or its `WITH`, reads as the
172            // whole compound's and the statement doesn't parse. Where a
173            // dialect has no parentheses for that (SQLite), a derived table
174            // says the same thing — and saying it unconditionally is what
175            // keeps a clause added later from slipping through.
176            if D::PARENTHESIZED_SET_OP_BRANCHES {
177                sink.ch('(');
178                sql.splice_into(sink);
179                sink.ch(')');
180            } else {
181                sink.text("SELECT * FROM (");
182                sql.splice_into(sink);
183                sink.ch(')');
184            }
185        };
186
187        // The chain is a left fold, and SQL's own precedence is not:
188        // `INTERSECT` binds tighter than `UNION`/`EXCEPT`, so flat text
189        // would reassociate `a.union(&b).intersect(&c)` into
190        // `A UNION (B INTERSECT C)` on Postgres — and, since SQLite reads
191        // compound operators left to right, would mean different things in
192        // the two dialects this crate executes. Parenthesising the
193        // accumulator wherever the operator changes says the fold outright,
194        // without encoding any dialect's precedence table.
195        let changes = self
196            .rest
197            .windows(2)
198            .filter(|pair| pair[0].0 != pair[1].0)
199            .count();
200        for _ in 0..changes {
201            if D::PARENTHESIZED_SET_OP_BRANCHES {
202                sink.ch('(');
203            } else {
204                sink.text("SELECT * FROM (");
205            }
206        }
207
208        branch(sink, &self.first);
209        for (i, (kind, part)) in self.rest.iter().enumerate() {
210            if i > 0 && self.rest[i - 1].0 != *kind {
211                sink.ch(')');
212            }
213            sink.text(kind.keyword());
214            branch(sink, part);
215        }
216    }
217
218    fn render_ordering(&self, sink: &mut QuerySink<D>) {
219        if !self.order_by.is_empty() {
220            sink.text(" ORDER BY ");
221            for (i, (position, dir)) in self.order_by.iter().enumerate() {
222                if i > 0 {
223                    sink.text(", ");
224                }
225                sink.text(&position.to_string());
226                sink.text(crate::render::dir_keyword(*dir));
227            }
228        }
229        crate::select::render_limit_offset::<D>(sink, self.limit.as_ref(), self.offset.as_ref());
230    }
231}
232
233impl<D: Dialect, Scope, Sel> Select<D, Scope, Sel> {
234    /// Starts a `UNION` chain — see `SetOp`'s doc comment for why the two
235    /// branches only need matching `Selection::Output`, not matching
236    /// `Scope`.
237    pub fn union<ScopeB, SelB, IdxA, IdxB>(
238        &self,
239        other: &Select<D, ScopeB, SelB>,
240    ) -> SetOp<D, Sel::Output>
241    where
242        Sel: Selection<Scope, IdxA>,
243        SelB: Selection<ScopeB, IdxB>,
244        SelB::Output: SameShape<Sel::Output>,
245    {
246        SetOp::new(self.fragment::<IdxA>()).union(other)
247    }
248
249    pub fn union_all<ScopeB, SelB, IdxA, IdxB>(
250        &self,
251        other: &Select<D, ScopeB, SelB>,
252    ) -> SetOp<D, Sel::Output>
253    where
254        Sel: Selection<Scope, IdxA>,
255        SelB: Selection<ScopeB, IdxB>,
256        SelB::Output: SameShape<Sel::Output>,
257    {
258        SetOp::new(self.fragment::<IdxA>()).union_all(other)
259    }
260
261    pub fn intersect<ScopeB, SelB, IdxA, IdxB>(
262        &self,
263        other: &Select<D, ScopeB, SelB>,
264    ) -> SetOp<D, Sel::Output>
265    where
266        Sel: Selection<Scope, IdxA>,
267        SelB: Selection<ScopeB, IdxB>,
268        SelB::Output: SameShape<Sel::Output>,
269    {
270        SetOp::new(self.fragment::<IdxA>()).intersect(other)
271    }
272
273    pub fn except<ScopeB, SelB, IdxA, IdxB>(
274        &self,
275        other: &Select<D, ScopeB, SelB>,
276    ) -> SetOp<D, Sel::Output>
277    where
278        Sel: Selection<Scope, IdxA>,
279        SelB: Selection<ScopeB, IdxB>,
280        SelB::Output: SameShape<Sel::Output>,
281    {
282        SetOp::new(self.fragment::<IdxA>()).except(other)
283    }
284}