use rupg_common::{Error, Result, SqlState};
use rupg_sql::nodes::{
CTEMaterialize, CommonTableExpr, JoinType, List, Node, RangeVar, SelectStmt, SetOperation,
WithClause,
};
use rupg_types::oid;
use crate::Analyzer;
use crate::coerce::AtOpt;
use crate::expr::Expr;
use crate::from::{Column, Item, Relation};
use crate::searchcycle::{self, CycleClause, SearchClause};
use crate::select::Query;
use crate::sort::{list_location, raw_location};
use crate::typename::place;
use crate::types;
use crate::walk::{self, Walker};
#[derive(Clone, Debug, PartialEq)]
pub struct Cte {
pub name: String,
pub aliases: Vec<String>,
pub columns: Vec<Column>,
pub query: Box<Query>,
pub recursive: bool,
pub materialized: Option<bool>,
pub search: Option<SearchClause>,
pub cycle: Option<CycleClause>,
pub write: Option<usize>,
}
#[derive(Clone, Debug, PartialEq)]
pub struct CteRef {
pub name: String,
pub work: Option<usize>,
pub write: Option<usize>,
pub origins: Vec<Option<(u32, i16)>>,
pub transition: bool,
}
#[derive(Clone, Debug)]
pub(crate) struct CteItem {
name: String,
location: Option<usize>,
aliases: Vec<String>,
recursive: bool,
materialized: Option<bool>,
visible: bool,
columns: Option<Vec<Column>>,
query: Option<Box<Query>>,
written: Option<Box<Query>>,
work: Option<usize>,
search: Option<SearchClause>,
cycle: Option<CycleClause>,
write: Option<usize>,
returning: bool,
origins: Vec<Option<(u32, i16)>>,
}
impl CteItem {
fn new(cte: &CommonTableExpr, recursive: bool, visible: bool) -> CteItem {
let materialized = match cte.ctematerialized {
CTEMaterialize::CTEMaterializeAlways => Some(true),
CTEMaterialize::CTEMaterializeNever => Some(false),
_ => None,
};
CteItem {
name: cte_name(cte).to_string(),
location: place(cte.location),
aliases: strings(&cte.aliascolnames),
recursive,
materialized,
visible,
columns: None,
query: None,
written: None,
work: None,
search: None,
cycle: None,
write: None,
returning: false,
origins: Vec::new(),
}
}
fn reference_columns(&self) -> Option<Vec<Column>> {
let mut columns = self.columns.clone()?;
columns.extend(searchcycle::extra_columns(self.search.as_ref(), self.cycle.as_ref()));
Some(columns)
}
fn into_cte(self) -> Result<Cte> {
let (Some(columns), Some(query)) = (self.reference_columns(), self.query) else {
return Err(Error::internal("a query of WITH with no analysis"));
};
Ok(Cte {
name: self.name,
aliases: self.aliases,
columns,
query: self.written.unwrap_or(query),
recursive: self.recursive,
materialized: self.materialized,
search: self.search,
cycle: self.cycle,
write: self.write,
})
}
fn columns(&self, shown: Vec<(String, u32, i32)>) -> Result<Vec<Column>> {
let mut columns = Vec::with_capacity(shown.len());
for (i, (name, ty, typmod)) in shown.into_iter().enumerate() {
let name = self.aliases.get(i).cloned().unwrap_or(name);
let (ty, typmod) =
if self.recursive && ty == oid::UNKNOWN { (oid::TEXT, -1) } else { (ty, typmod) };
columns.push(Column { name, ty, typmod, not_null: false });
}
if columns.len() < self.aliases.len() {
return Err(Error::new(
SqlState::INVALID_COLUMN_REFERENCE,
format!(
"WITH query \"{}\" has {} columns available but {} columns specified",
self.name,
columns.len(),
self.aliases.len()
),
)
.at_opt(self.location));
}
Ok(columns)
}
}
fn cte_name(cte: &CommonTableExpr) -> &str {
cte.ctename.as_deref().unwrap_or_default()
}
fn strings(list: &List) -> Vec<String> {
list.iter()
.filter_map(|n| match n {
Some(Node::String(s)) => Some(s.to_string()),
_ => None,
})
.collect()
}
fn ctes(with: &WithClause) -> Vec<&CommonTableExpr> {
with.ctes
.iter()
.filter_map(|n| match n {
Some(Node::CommonTableExpr(cte)) => Some(&**cte),
_ => None,
})
.collect()
}
fn unqualified(rv: &RangeVar) -> Option<&str> {
match rv.schemaname {
Some(_) => None,
None => Some(rv.relname.as_deref().unwrap_or_default()),
}
}
#[derive(Default)]
struct Inner(Vec<Vec<String>>);
impl Inner {
fn hides(&self, name: &str) -> bool {
self.0.iter().flatten().any(|n| n == name)
}
}
struct Deps<'a> {
names: &'a [&'a str],
current: usize,
inner: Inner,
depends: Vec<usize>,
recursive: bool,
}
impl Deps<'_> {
fn inner_with(&mut self, s: &SelectStmt, with: &WithClause) -> bool {
let list = ctes(with);
if with.recursive {
self.inner.0.push(list.iter().map(|c| cte_name(c).to_string()).collect());
for cte in &list {
walk::one(self, cte.ctequery.as_ref());
}
} else {
self.inner.0.push(Vec::new());
for cte in &list {
walk::one(self, cte.ctequery.as_ref());
if let Some(level) = self.inner.0.last_mut() {
level.push(cte_name(cte).to_string());
}
}
}
walk::select(self, s);
self.inner.0.pop();
false
}
}
impl Walker for Deps<'_> {
fn node(&mut self, node: &Node) -> bool {
match node {
Node::RangeVar(rv) => {
if let Some(name) = unqualified(rv)
&& !self.inner.hides(name)
&& let Some(i) = self.names.iter().position(|n| *n == name)
{
if i == self.current {
self.recursive = true;
} else if !self.depends.contains(&i) {
self.depends.push(i);
}
}
false
}
Node::SelectStmt(s) => self.select(s),
Node::WithClause(_) => false,
_ => walk::walk(self, node),
}
}
fn select(&mut self, s: &SelectStmt) -> bool {
match s.withClause.as_deref() {
Some(with) => self.inner_with(s, with),
None => walk::select(self, s),
}
}
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum Context {
Ok,
NonRecursiveTerm,
SubLink,
OuterJoin,
Intersect,
Except,
}
impl Context {
fn message(self, name: &str) -> String {
let place = match self {
Context::Ok => "",
Context::NonRecursiveTerm => "its non-recursive term",
Context::SubLink => "a subquery",
Context::OuterJoin => "an outer join",
Context::Intersect => "INTERSECT",
Context::Except => "EXCEPT",
};
format!("recursive reference to query \"{name}\" must not appear within {place}")
}
}
struct Recursion<'a> {
name: &'a str,
inner: Inner,
count: usize,
context: Context,
error: Option<Error>,
}
impl Recursion<'_> {
fn new(name: &str, context: Context) -> Recursion<'_> {
Recursion { name, inner: Inner::default(), count: 0, context, error: None }
}
fn fail(&mut self, error: Error) -> bool {
self.error = Some(error);
true
}
fn select_stmt(&mut self, s: &SelectStmt) -> bool {
let save = self.context;
if save != Context::Ok {
return walk::select(self, s);
}
let (left, right) = match s.op {
SetOperation::SETOP_INTERSECT if s.all => (Context::Intersect, Context::Intersect),
SetOperation::SETOP_INTERSECT => (save, save),
SetOperation::SETOP_EXCEPT if s.all => (Context::Except, Context::Except),
SetOperation::SETOP_EXCEPT => (save, Context::Except),
_ => return walk::select(self, s),
};
self.context = left;
if s.larg.as_deref().is_some_and(|l| self.select(l)) {
return true;
}
self.context = right;
if s.rarg.as_deref().is_some_and(|r| self.select(r)) {
return true;
}
self.context = save;
walk::list(self, &s.sortClause)
|| walk::one(self, s.limitOffset.as_ref())
|| walk::one(self, s.limitCount.as_ref())
|| walk::list(self, &s.lockingClause)
}
fn within(&mut self, context: Context, save: Context, node: Option<&Node>) -> bool {
self.context = context;
let stop = walk::one(self, node);
self.context = save;
stop
}
}
impl Walker for Recursion<'_> {
fn node(&mut self, node: &Node) -> bool {
let save = self.context;
let outer = if save == Context::Ok { Context::OuterJoin } else { save };
match node {
Node::RangeVar(rv) => {
let Some(name) = unqualified(rv) else { return false };
if self.inner.hides(name) || name != self.name {
return false;
}
let at = place(rv.location);
if self.context != Context::Ok {
let message = self.context.message(self.name);
return self.fail(Error::new(SqlState::INVALID_RECURSION, message).at_opt(at));
}
self.count += 1;
if self.count > 1 {
let message = format!(
"recursive reference to query \"{}\" must not appear more than once",
self.name
);
return self.fail(Error::new(SqlState::INVALID_RECURSION, message).at_opt(at));
}
false
}
Node::SelectStmt(s) => self.select(s),
Node::WithClause(_) => false,
Node::JoinExpr(j) => {
let (larg, rarg) = (j.larg.as_ref(), j.rarg.as_ref());
let sides = match j.jointype {
JoinType::JOIN_LEFT => walk::one(self, larg) || self.within(outer, save, rarg),
JoinType::JOIN_FULL => {
self.within(outer, save, larg) || self.within(outer, save, rarg)
}
JoinType::JOIN_RIGHT => self.within(outer, save, larg) || walk::one(self, rarg),
_ => walk::one(self, larg) || walk::one(self, rarg),
};
sides || walk::one(self, j.quals.as_ref())
}
Node::SubLink(s) => {
self.within(Context::SubLink, save, s.subselect.as_ref())
|| walk::one(self, s.testexpr.as_ref())
}
_ => walk::walk(self, node),
}
}
fn select(&mut self, s: &SelectStmt) -> bool {
let Some(with) = s.withClause.as_deref() else { return self.select_stmt(s) };
let list = ctes(with);
if with.recursive {
self.inner.0.push(list.iter().map(|c| cte_name(c).to_string()).collect());
for cte in &list {
if walk::one(self, cte.ctequery.as_ref()) {
return true;
}
}
} else {
self.inner.0.push(Vec::new());
for cte in &list {
if walk::one(self, cte.ctequery.as_ref()) {
return true;
}
if let Some(level) = self.inner.0.last_mut() {
level.push(cte_name(cte).to_string());
}
}
}
let stop = self.select_stmt(s);
self.inner.0.pop();
stop
}
}
fn check_walk(
name: &str,
context: Context,
f: impl FnOnce(&mut Recursion<'_>) -> bool,
) -> Result<usize> {
let mut walker = Recursion::new(name, context);
f(&mut walker);
match walker.error {
Some(error) => Err(error),
None => Ok(walker.count),
}
}
fn check_recursion(cte: &CommonTableExpr) -> Result<()> {
let name = cte_name(cte);
let Some(Node::SelectStmt(stmt)) = &cte.ctequery else {
return Err(Error::new(
SqlState::INVALID_RECURSION,
format!("recursive query \"{name}\" must not contain data-modifying statements"),
)
.at_opt(place(cte.location)));
};
if stmt.op != SetOperation::SETOP_UNION {
return Err(Error::new(
SqlState::INVALID_RECURSION,
format!(
"recursive query \"{name}\" does not have the form non-recursive-term UNION [ALL] recursive-term"
),
)
.at_opt(place(cte.location)));
}
if let Some(with) = stmt.withClause.as_deref() {
check_walk(name, Context::SubLink, |w| walk::list(w, &with.ctes))?;
}
let not_implemented = |what: &str, at: Option<usize>| {
Err(Error::new(
SqlState::FEATURE_NOT_SUPPORTED,
format!("{what} in a recursive query is not implemented"),
)
.at_opt(at))
};
if !stmt.sortClause.is_empty() {
return not_implemented("ORDER BY", list_location(&stmt.sortClause));
}
if stmt.limitOffset.is_some() {
return not_implemented("OFFSET", raw_location(stmt.limitOffset.as_ref()));
}
if stmt.limitCount.is_some() {
return not_implemented("LIMIT", raw_location(stmt.limitCount.as_ref()));
}
if !stmt.lockingClause.is_empty() {
return not_implemented("FOR UPDATE/SHARE", None);
}
let larg = stmt.larg.as_deref();
check_walk(name, Context::NonRecursiveTerm, |w| larg.is_some_and(|l| w.select(l)))?;
let rarg = stmt.rarg.as_deref();
let count = check_walk(name, Context::Ok, |w| rarg.is_some_and(|r| w.select(r)))?;
if count != 1 {
return Err(Error::internal("missing recursive reference"));
}
Ok(())
}
fn raise_outer(query: &mut Query, levels: usize) {
query.each_expr_mut(0, &mut |expr, depth| expr.raise_from(levels, depth + 1));
}
impl Analyzer<'_> {
pub(crate) fn transform_with(&mut self, with: &WithClause) -> Result<()> {
let list = ctes(with);
for (i, cte) in list.iter().enumerate() {
if let Some(again) = list[i + 1..].iter().find(|c| cte_name(c) == cte_name(cte)) {
return Err(Error::new(
SqlState::DUPLICATE_ALIAS,
format!("WITH query name \"{}\" specified more than once", cte_name(again)),
)
.at_opt(place(again.location)));
}
if !matches!(cte.ctequery, Some(Node::SelectStmt(_))) && !self.outer.is_empty() {
return Err(Error::new(
SqlState::FEATURE_NOT_SUPPORTED,
"WITH clause containing a data-modifying statement must be at the top level",
)
.at_opt(place(cte.location)));
}
}
let base = self.scope.ctes.len();
if !with.recursive {
self.scope.ctes.extend(list.iter().map(|c| CteItem::new(c, false, false)));
for (i, cte) in list.iter().enumerate() {
self.analyze_cte(base + i, cte)?;
self.scope.ctes[base + i].visible = true;
}
return Ok(());
}
let names: Vec<&str> = list.iter().map(|c| cte_name(c)).collect();
let mut order = Vec::with_capacity(list.len());
for (i, cte) in list.iter().enumerate() {
let mut deps = Deps {
names: &names,
current: i,
inner: Inner::default(),
depends: Vec::new(),
recursive: false,
};
walk::one(&mut deps, cte.ctequery.as_ref());
order.push((i, deps.depends, deps.recursive));
}
for i in 0..order.len() {
let Some(j) = (i..order.len()).find(|&j| order[j].1.is_empty()) else {
return Err(Error::new(
SqlState::FEATURE_NOT_SUPPORTED,
"mutual recursion between WITH items is not implemented",
)
.at_opt(place(list[order[i].0].location)));
};
order.swap(i, j);
let id = order[i].0;
for later in &mut order[i + 1..] {
later.1.retain(|d| *d != id);
}
}
for (id, _, recursive) in &order {
if *recursive {
check_recursion(list[*id])?;
}
}
self.scope
.ctes
.extend(order.iter().map(|(id, _, rec)| CteItem::new(list[*id], *rec, true)));
for (k, (id, _, _)) in order.iter().enumerate() {
self.analyze_cte(base + k, list[*id])?;
}
Ok(())
}
fn analyze_cte(&mut self, slot: usize, cte: &CommonTableExpr) -> Result<()> {
let stmt = match &cte.ctequery {
Some(Node::SelectStmt(stmt)) => stmt,
Some(node) => return self.analyze_write_cte(slot, cte, node),
None => return Err(Error::internal("a query of WITH with no statement")),
};
self.scope.ctes[slot].search = cte.search_clause.as_deref().map(searchcycle::search_clause);
if let Some(raw) = cte.cycle_clause.as_deref() {
self.scope.ctes[slot].cycle = Some(self.cycle_clause(raw)?);
}
if self.scope.ctes[slot].recursive {
self.scope.ctes[slot].work = Some(self.works);
self.works += 1;
}
self.parent_cte = Some(slot);
let query = self.sub_select(stmt, true);
self.parent_cte = None;
let query = query?;
let item = &self.scope.ctes[slot];
let shown = query.targets.iter().filter(|t| !t.junk);
if item.recursive {
let Some(columns) = &item.columns else {
return Err(Error::internal("a recursive query with no columns"));
};
for (i, (target, column)) in shown.zip(columns).enumerate() {
let expr = &target.expr;
if expr.ty != column.ty || expr.typmod != column.typmod {
return Err(Error::new(
SqlState::DATATYPE_MISMATCH,
format!(
"recursive query \"{}\" column {} has type {} in non-recursive term but type {} overall",
item.name,
i + 1,
types::with_typmod(column.ty, column.typmod),
types::with_typmod(expr.ty, expr.typmod)
),
)
.with_hint("Cast the output of the non-recursive term to the correct type.")
.at_opt(expr.place()));
}
}
} else {
let shown = shown.map(|t| (t.name.clone(), t.expr.ty, t.expr.typmod)).collect();
let columns = item.columns(shown)?;
self.scope.ctes[slot].columns = Some(columns);
}
let item = &mut self.scope.ctes[slot];
let (search, cycle) = (item.search.as_ref(), item.cycle.as_ref());
if search.is_some() || cycle.is_some() {
let columns = item.columns.as_deref().unwrap_or_default();
searchcycle::check(columns, item.recursive, item.location, search, cycle, &query)?;
let rewritten =
searchcycle::rewrite(&item.name, columns, item.work, search, cycle, &query)?;
item.written = Some(Box::new(query));
item.query = Some(Box::new(rewritten));
return Ok(());
}
item.query = Some(Box::new(query));
Ok(())
}
fn analyze_write_cte(&mut self, slot: usize, cte: &CommonTableExpr, node: &Node) -> Result<()> {
let query = self.sub_level(true, |analyzer| match node {
Node::InsertStmt(s) => analyzer.insert(s),
Node::UpdateStmt(s) => analyzer.update(s),
Node::DeleteStmt(s) => analyzer.delete(s),
Node::MergeStmt(s) => analyzer.merge(s),
_ => Err(Error::internal("a query of WITH that is not a statement of a query")),
})?;
let returning = query.write.as_ref().and_then(|w| w.returning.as_deref());
let shown: Vec<&crate::select::Target> =
returning.map(|r| r.targets.iter().filter(|t| !t.junk).collect()).unwrap_or_default();
let item = &self.scope.ctes[slot];
let columns = item
.columns(shown.iter().map(|t| (t.name.clone(), t.expr.ty, t.expr.typmod)).collect())?;
if cte.search_clause.is_some() || cte.cycle_clause.is_some() {
return Err(Error::new(SqlState::SYNTAX_ERROR, "WITH query is not recursive")
.at_opt(item.location));
}
let item = &mut self.scope.ctes[slot];
item.origins = shown.iter().map(|t| t.origin).collect();
item.returning = returning.is_some();
item.columns = Some(columns);
item.query = Some(Box::new(query));
item.write = Some(self.writes);
self.writes += 1;
Ok(())
}
pub(crate) fn recursive_parent(&self, parent: Option<usize>) -> Option<&CteItem> {
let item = self.outer.last()?.ctes.get(parent?)?;
item.recursive.then_some(item)
}
pub(crate) fn work_table(&self, parent: Option<usize>) -> Option<usize> {
self.recursive_parent(parent).and_then(|item| item.work)
}
pub(crate) fn recursive_columns(
&mut self,
parent: usize,
leftmost: usize,
exprs: &[Expr],
) -> Result<()> {
let Some(query) = self.scope.relations[leftmost].subquery.as_deref() else {
return Err(Error::internal("a set operation with no leftmost SELECT"));
};
let names = query.targets.iter().filter(|t| !t.junk).map(|t| t.name.clone());
let shown = names.zip(exprs).map(|(n, e)| (n, e.ty, e.typmod)).collect();
let Some(item) = self.outer.last_mut().and_then(|o| o.ctes.get_mut(parent)) else {
return Err(Error::internal("a recursive query that is not in the scope"));
};
item.columns = Some(item.columns(shown)?);
Ok(())
}
pub(crate) fn cte_reference(&mut self, rv: &RangeVar) -> Result<Option<Item>> {
let Some(name) = unqualified(rv) else { return Ok(None) };
let found = self.levels().find_map(|(level, scope)| {
scope.ctes.iter().find(|c| c.visible && c.name == name).map(|c| (level, c))
});
let Some((level, item)) = found else { return Ok(None) };
let Some(columns) = item.reference_columns() else {
return Err(Error::internal("a reference to a query of WITH with no columns"));
};
if let Some(write) = item.write {
if !item.returning {
return Err(Error::new(
SqlState::FEATURE_NOT_SUPPORTED,
format!("WITH query \"{name}\" does not have a RETURNING clause"),
)
.at_opt(place(rv.location)));
}
let mut relation = Relation::new(0, columns, None, None, None);
relation.cte = Some(CteRef {
name: name.to_string(),
work: None,
write: Some(write),
origins: item.origins.clone(),
transition: false,
});
let alias = rv.alias.as_deref();
let refname = alias.and_then(|a| a.aliasname.as_deref()).unwrap_or(name).to_string();
return self.add_relation(relation, refname, alias).map(Some);
}
let available = item.columns.as_ref().map_or(0, Vec::len);
let extra = columns.len() - available;
let mut relation = match &item.query {
Some(query) => {
let mut query = (**query).clone();
if level > 0 {
raise_outer(&mut query, level);
}
Relation::new(0, columns, Some(Box::new(query)), None, None)
}
None => Relation::new(0, columns, None, None, None),
};
let work = if item.query.is_none() { item.work } else { None };
relation.cte = Some(CteRef {
name: name.to_string(),
work,
write: None,
origins: Vec::new(),
transition: false,
});
let alias = rv.alias.as_deref();
let refname = alias.and_then(|a| a.aliasname.as_deref()).unwrap_or(name).to_string();
let given = alias.map_or(0, |a| a.colnames.len());
if extra > 0 && given > available {
return Err(Error::new(
SqlState::INVALID_COLUMN_REFERENCE,
format!(
"table \"{refname}\" has {available} columns available but {given} columns specified"
),
));
}
let hidden = if level > 0 { extra } else { 0 };
let mut item = self.add_relation(relation, refname, alias)?;
item.hide_last(hidden);
Ok(Some(item))
}
pub(crate) fn transition_reference(&mut self, rv: &RangeVar) -> Result<Option<Item>> {
let Some(name) = unqualified(rv) else { return Ok(None) };
let Some(table) = self.env.transition_table(name) else { return Ok(None) };
let Some(rel) = self.env.catalog().and_then(|c| c.relation(table)) else {
return Ok(None);
};
let columns = rel
.columns
.iter()
.map(|c| Column {
name: if c.dropped { String::new() } else { c.name.clone() },
ty: if c.dropped { 0 } else { c.ty },
typmod: c.typmod,
not_null: false,
})
.collect();
let mut relation = Relation::new(0, columns, None, None, None);
relation.cte = Some(CteRef {
name: name.to_string(),
work: None,
write: None,
origins: Vec::new(),
transition: true,
});
let alias = rv.alias.as_deref();
let refname = alias.and_then(|a| a.aliasname.as_deref()).unwrap_or(name).to_string();
self.add_relation(relation, refname, alias).map(Some)
}
pub(crate) fn future_cte(&self, name: &str) -> bool {
self.levels().any(|(_, scope)| scope.ctes.iter().any(|c| !c.visible && c.name == name))
}
pub(crate) fn write_with(&mut self, with: Option<&WithClause>) -> Result<()> {
match with {
Some(with) => self.transform_with(with),
None => Ok(()),
}
}
pub(crate) fn end_write_with(
&mut self,
query: &mut Query,
with: Option<&WithClause>,
) -> Result<()> {
self.finish_auto_view(query);
query.recursive = with.is_some_and(|w| w.recursive);
query.ctes = self.take_ctes()?;
Ok(())
}
pub(crate) fn take_ctes(&mut self) -> Result<Vec<Cte>> {
std::mem::take(&mut self.scope.ctes).into_iter().map(CteItem::into_cte).collect()
}
pub(crate) fn check_recursive_aggs<'e>(
&self,
exprs: impl Iterator<Item = &'e Expr>,
) -> Result<()> {
let self_reference =
self.scope.relations.iter().any(|r| r.cte.as_ref().is_some_and(|c| c.work.is_some()));
if !self.has_aggs || !self_reference {
return Ok(());
}
let at = exprs.into_iter().find_map(Expr::first_agg).and_then(Expr::place);
Err(Error::new(
SqlState::INVALID_RECURSION,
"aggregate functions are not allowed in a recursive query's recursive term",
)
.at_opt(at))
}
}