1use 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
42pub 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 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
74impl<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 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 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 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 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 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 fn render_branches(&self, sink: &mut QuerySink<D>) {
173 let branch = |sink: &mut QuerySink<D>, sql: &Fragment| {
174 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 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 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}