1use rupg_common::{Error, Result, SqlState};
4use rupg_sql::nodes::{LimitOption, List, Node, SelectStmt, SetOperation};
5use rupg_types::oid;
6
7use crate::Analyzer;
8use crate::agg::Kind;
9use crate::coerce::Context;
10use crate::colname::figure_colname;
11use crate::expr::{Expr, ExprKind};
12use crate::from::{FromItem, Relation};
13use crate::sort::SortGroup;
14use crate::window::WindowClause;
15
16#[derive(Clone, Debug, Default, PartialEq)]
18pub struct Query {
19 pub relations: Vec<Relation>,
21 pub from: Vec<FromItem>,
23 pub targets: Vec<Target>,
25 pub filter: Option<Expr>,
27 pub grouped: bool,
29 pub group: Vec<SortGroup>,
31 pub grouping_sets: Vec<crate::GroupingSet>,
33 pub group_distinct: bool,
35 pub having: Option<Expr>,
37 pub sort: Vec<SortGroup>,
39 pub distinct: Vec<SortGroup>,
41 pub distinct_on: bool,
43 pub offset: Option<Expr>,
45 pub limit: Option<Expr>,
47 pub with_ties: bool,
49 pub params: Vec<u32>,
51 pub notices: Vec<Error>,
53 pub set_op: Option<crate::SetOp>,
55 pub windows: Vec<WindowClause>,
57 pub ctes: Vec<crate::Cte>,
59 pub recursive: bool,
61 pub row_marks: Vec<crate::RowMark>,
63 pub write: Option<crate::Write>,
65 pub out_args: Vec<crate::OutArg>,
67 pub pulled: usize,
69 pub pushed: bool,
71 pub barrier: Vec<usize>,
73 pub minmax: Option<Box<MinMax>>,
75 pub constraint_deps: Vec<u32>,
77 pub shown: Option<Box<crate::Shown>>,
79 pub perms: Vec<crate::PermInfo>,
81 pub funcs: Vec<u32>,
83}
84
85#[derive(Clone, Debug, PartialEq)]
87pub struct MinMax {
88 pub aggs: Vec<[Query; 2]>,
90 pub targets: Vec<Expr>,
92 pub having: Option<Expr>,
94}
95
96#[derive(Clone, Debug, PartialEq)]
98pub struct Target {
99 pub name: String,
101 pub expr: Expr,
102 pub origin: Option<(u32, i16)>,
104 pub junk: bool,
106}
107
108impl Query {
109 pub fn exprs(&self) -> Vec<&Expr> {
111 fn joins<'a>(item: &'a FromItem, out: &mut Vec<&'a Expr>) {
112 if let FromItem::Join(j) = item {
113 joins(&j.left, out);
114 joins(&j.right, out);
115 out.extend(&j.on);
116 }
117 }
118 let mut out: Vec<&Expr> = self.targets.iter().map(|t| &t.expr).collect();
119 for item in &self.from {
120 joins(item, &mut out);
121 }
122 out.extend(self.filter.iter().chain(&self.having).chain(&self.offset).chain(&self.limit));
123 for window in &self.windows {
124 out.extend(window.frame.start.iter().chain(&window.frame.end));
125 }
126 for relation in &self.relations {
127 if let Some(function) = &relation.function {
128 out.extend(function.calls.iter().map(|(call, _)| call));
129 }
130 out.extend(relation.values.iter().flatten().flatten());
131 if let Some(sample) = &relation.tablesample {
132 out.extend(sample.args.iter().chain(&sample.repeatable));
133 }
134 }
135 out
136 }
137
138 pub fn exprs_mut(&mut self) -> Vec<&mut Expr> {
140 fn joins<'a>(item: &'a mut FromItem, out: &mut Vec<&'a mut Expr>) {
141 if let FromItem::Join(j) = item {
142 let j = &mut **j;
143 joins(&mut j.left, out);
144 joins(&mut j.right, out);
145 out.extend(&mut j.on);
146 }
147 }
148 let mut out: Vec<&mut Expr> = self.targets.iter_mut().map(|t| &mut t.expr).collect();
149 for item in &mut self.from {
150 joins(item, &mut out);
151 }
152 out.extend(
153 self.filter
154 .iter_mut()
155 .chain(&mut self.having)
156 .chain(&mut self.offset)
157 .chain(&mut self.limit),
158 );
159 for window in &mut self.windows {
160 out.extend(window.frame.start.iter_mut().chain(&mut window.frame.end));
161 }
162 for relation in &mut self.relations {
163 if let Some(function) = &mut relation.function {
164 out.extend(function.calls.iter_mut().map(|(call, _)| call));
165 }
166 out.extend(relation.values.iter_mut().flatten().flatten());
167 if let Some(sample) = &mut relation.tablesample {
168 out.extend(sample.args.iter_mut().chain(&mut sample.repeatable));
169 }
170 }
171 out
172 }
173
174 pub fn has_target_srfs(&self) -> bool {
176 self.targets.iter().any(|t| t.expr.first_set_call().is_some())
177 }
178
179 pub fn all_exprs(&self, depth: usize) -> Vec<(&Expr, usize)> {
181 let mut out: Vec<(&Expr, usize)> = self.exprs().into_iter().map(|e| (e, depth)).collect();
182 for relation in &self.relations {
183 if let Some(sub) = &relation.subquery {
184 out.extend(sub.all_exprs(depth + 1));
185 }
186 }
187 out
188 }
189
190 pub fn each_expr_mut(&mut self, depth: usize, f: &mut impl FnMut(&mut Expr, usize)) {
192 for expr in self.exprs_mut() {
193 f(expr, depth);
194 }
195 for relation in &mut self.relations {
196 if let Some(sub) = &mut relation.subquery {
197 sub.each_expr_mut(depth + 1, f);
198 }
199 }
200 }
201}
202
203fn not_yet(what: &str) -> Error {
205 Error::new(SqlState::FEATURE_NOT_SUPPORTED, format!("{what} is not supported yet"))
206}
207
208impl Analyzer<'_> {
209 pub(crate) fn select(&mut self, s: &SelectStmt) -> Result<Query> {
211 let keep_unknowns = std::mem::take(&mut self.keep_unknowns);
213 let parent = self.parent_cte.take();
214 if let Some(with) = s.withClause.as_deref() {
215 self.transform_with(with)?;
216 }
217 let mut query = if s.op != SetOperation::SETOP_NONE {
218 self.set_operation(s, parent)?
219 } else if !s.valuesLists.is_empty() {
220 self.values_clause(s)?
221 } else {
222 self.simple_select(s, keep_unknowns)?
223 };
224 query.recursive = s.withClause.as_deref().is_some_and(|w| w.recursive);
225 query.ctes = self.take_ctes()?;
226 Ok(query)
227 }
228
229 pub(crate) fn target_list(&mut self, list: &List, kind: Kind) -> Result<Vec<Target>> {
231 let mut targets = Vec::with_capacity(list.len());
232 for target in list {
233 let Some(Node::ResTarget(rt)) = target else {
234 return Err(Error::internal("a target that is not ResTarget"));
235 };
236 if let Some(Node::ColumnRef(c)) = &rt.val
237 && let Some(columns) = self.expand_star(c)?
238 {
239 for (name, expr) in columns {
240 targets.push(Target { name, expr, origin: None, junk: false });
241 }
242 continue;
243 }
244 if let Some(Node::A_Indirection(ind)) = &rt.val
245 && let Some(fields) = self.with_kind(kind, |a| a.expand_indirection_star(ind))?
246 {
247 for (name, expr) in fields {
248 targets.push(Target { name, expr, origin: None, junk: false });
249 }
250 continue;
251 }
252 let expr = self.with_kind(kind, |a| a.transform(rt.val.as_ref()))?;
253 let name = match &rt.name {
254 Some(name) => name.to_string(),
255 None => figure_colname(rt.val.as_ref()),
256 };
257 targets.push(Target { name, expr, origin: None, junk: false });
258 }
259 Ok(targets)
260 }
261
262 pub(crate) fn mark_origins(&self, targets: &mut [Target]) {
264 for target in targets {
265 if let ExprKind::Var(var) = &target.expr.kind {
266 let relation = &self.level(var.levels_up).relations[var.relation];
267 let written = relation.cte.as_ref().filter(|c| c.write.is_some());
268 target.origin = match (&relation.subquery, written) {
269 (None, Some(cte)) => usize::try_from(var.attnum - 1)
270 .ok()
271 .and_then(|i| cte.origins.get(i).copied().flatten()),
272 (None, None)
273 if relation.function.is_some()
274 || relation.values.is_some()
275 || relation.cte.is_some() =>
276 {
277 None
278 }
279 (None, None) => Some((relation.oid, var.attnum)),
280 (Some(_), _) if relation.oid != 0 => Some((relation.oid, var.attnum)),
281 (Some(sub), _) => usize::try_from(var.attnum - 1)
282 .ok()
283 .and_then(|i| sub.targets.iter().filter(|t| !t.junk).nth(i))
284 .and_then(|t| t.origin),
285 };
286 }
287 }
288 }
289
290 fn simple_select(&mut self, s: &SelectStmt, keep_unknowns: bool) -> Result<Query> {
292 if s.intoClause.is_some() {
293 return Err(not_yet("SELECT INTO"));
294 }
295 let defs = s.windowClause.iter().flatten().filter_map(|node| match node {
297 Node::WindowDef(def) => Some((**def).clone()),
298 _ => None,
299 });
300 let windowdefs = std::mem::replace(&mut self.windowdefs, defs.collect());
301 let has_windows = std::mem::replace(&mut self.has_windows, false);
302 let result = self.plain_select(s, keep_unknowns);
303 self.windowdefs = windowdefs;
304 self.has_windows = has_windows;
305 result
306 }
307
308 fn plain_select(&mut self, s: &SelectStmt, keep_unknowns: bool) -> Result<Query> {
310 let assign = self.assign.take();
312 let from = self.transform_from(&s.fromClause)?;
313 let mut targets = self.target_list(&s.targetList, Kind::Select)?;
314 if let Some(assign) = assign {
315 self.assign_target(&mut targets, assign)?;
316 }
317 let filter = match &s.whereClause {
318 Some(node) => {
319 let expr = self.with_kind(Kind::Where, |a| a.transform(Some(node)))?;
320 Some(self.coerce_to_boolean(expr, "WHERE")?)
321 }
322 None => None,
323 };
324 let mut having = match &s.havingClause {
325 Some(node) => {
326 let expr = self.with_kind(Kind::Having, |a| a.transform(Some(node)))?;
327 Some(self.coerce_to_boolean(expr, "HAVING")?)
328 }
329 None => None,
330 };
331 self.mark_origins(&mut targets);
332 let sort = self.sort_clause(&s.sortClause, &mut targets)?;
333 let (group, mut grouping_sets) = self.group_clause(&s.groupClause, &mut targets, &sort)?;
334 let distinct_on = matches!(s.distinctClause.first(), Some(Some(_)));
336 let distinct = if s.distinctClause.is_empty() {
337 Vec::new()
338 } else if distinct_on {
339 self.distinct_on_clause(&s.distinctClause, &mut targets, &sort)?
340 } else {
341 self.distinct_clause(&mut targets, &sort, false)?
342 };
343 let with_ties = s.limitOption == LimitOption::LIMIT_OPTION_WITH_TIES;
344 let offset = self.with_kind(Kind::Offset, |a| {
345 a.limit_clause(s.limitOffset.as_ref(), "OFFSET", with_ties)
346 })?;
347 let limit = self.with_kind(Kind::Limit, |a| {
348 a.limit_clause(s.limitCount.as_ref(), "LIMIT", with_ties)
349 })?;
350 let windows = self.window_clauses(&mut targets)?;
351 for target in targets.iter_mut().filter(|_| !keep_unknowns) {
353 if target.expr.ty == oid::UNKNOWN {
354 let expr = std::mem::replace(&mut target.expr, Expr::new(ExprKind::CaseTest, 0));
355 target.expr = self.coerce(expr, oid::TEXT, -1, Context::Implicit, None)?;
356 }
357 }
358 let group_distinct = s.groupDistinct;
359 let grouped =
360 self.has_aggs || !group.is_empty() || !grouping_sets.is_empty() || having.is_some();
361 let mut constraint_deps = Vec::new();
362 if grouped {
363 constraint_deps = self.check_grouping(
364 &mut targets,
365 having.as_mut(),
366 &group,
367 &mut grouping_sets,
368 group_distinct,
369 )?;
370 }
371 let relations = std::mem::take(&mut self.scope.relations);
372 let mut query = Query {
373 relations,
374 from,
375 targets,
376 filter,
377 grouped,
378 group,
379 grouping_sets,
380 group_distinct,
381 having,
382 sort,
383 distinct,
384 distinct_on,
385 offset,
386 limit,
387 with_ties,
388 params: Vec::new(),
389 notices: Vec::new(),
390 set_op: None,
391 windows,
392 ctes: Vec::new(),
393 recursive: false,
394 row_marks: Vec::new(),
395 write: None,
396 out_args: Vec::new(),
397 pulled: 0,
398 pushed: false,
399 barrier: Vec::new(),
400 minmax: None,
401 constraint_deps,
402 shown: None,
403 perms: Vec::new(),
404 funcs: Vec::new(),
405 };
406 for clause in s.lockingClause.iter().flatten() {
407 if let Node::LockingClause(clause) = clause {
408 crate::locking::transform_locking_clause(&mut query, clause)?;
409 }
410 }
411 if !s.lockingClause.is_empty() {
412 self.mark_locked(&query);
413 }
414 Ok(query)
415 }
416}