1use 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
38pub 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 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
70impl<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 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 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 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 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 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 fn render_branches(&self, sink: &mut QuerySink<D>) {
169 let branch = |sink: &mut QuerySink<D>, sql: &Fragment| {
170 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 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 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}