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