use std::collections::BTreeSet;
use crate::identifier::{Identifier, QualifiedName};
use crate::ir::catalog::Catalog;
use crate::ir::column_type::ColumnType;
use crate::ir::index::IndexParent;
use crate::ir::view::ViewColumn;
use crate::parse::normalize_body::NormalizedBody;
use crate::plan::edges::{DepEdge, DepSource, NodeId};
#[derive(Debug, thiserror::Error)]
pub enum AstCanonError {
#[error("view {view}: failed to canonicalize body: {reason}")]
NormalizeFailed {
view: String,
reason: String,
},
#[error("view {view}: references {object} which is not declared in source")]
UnresolvedReference {
view: String,
object: String,
},
}
pub fn canonicalize_view_bodies(catalog: &mut Catalog) -> Result<(), AstCanonError> {
let known = KnownObjects::from_catalog(catalog);
for i in 0..catalog.views.len() {
let raw_body = catalog.views[i].raw_body.clone();
let qname = catalog.views[i].qname.clone();
let qname_str = qname.to_string();
let column_aliases: Vec<String> = catalog.views[i]
.columns
.iter()
.map(|c| c.name.as_str().to_string())
.collect();
let has_explicit_columns = !column_aliases.is_empty();
let effective_body = if has_explicit_columns {
apply_view_column_aliases(&raw_body, &column_aliases)
.unwrap_or_else(|| raw_body.clone())
} else {
raw_body.clone()
};
let normalized = NormalizedBody::from_sql(&effective_body).map_err(|e| {
AstCanonError::NormalizeFailed {
view: qname_str.clone(),
reason: e.to_string(),
}
})?;
let (deps, derived_columns) = walk_body_ast(&raw_body, &qname, &known)?;
catalog.views[i].body_canonical = normalized;
catalog.views[i].body_dependencies = deps;
if has_explicit_columns {
resolve_explicit_column_types(&mut catalog.views[i].columns);
} else {
catalog.views[i].columns = derived_columns;
}
}
for i in 0..catalog.materialized_views.len() {
let raw_body = catalog.materialized_views[i].raw_body.clone();
let qname = catalog.materialized_views[i].qname.clone();
let has_explicit_columns = !catalog.materialized_views[i].columns.is_empty();
let normalized =
NormalizedBody::from_sql(&raw_body).map_err(|e| AstCanonError::NormalizeFailed {
view: qname.to_string(),
reason: e.to_string(),
})?;
let (deps, derived_columns) = walk_body_ast(&raw_body, &qname, &known)?;
catalog.materialized_views[i].body_canonical = normalized;
catalog.materialized_views[i].body_dependencies = deps;
if has_explicit_columns {
resolve_explicit_column_types(&mut catalog.materialized_views[i].columns);
} else {
catalog.materialized_views[i].columns = derived_columns;
}
}
Ok(())
}
pub fn promote_mv_index_parents(catalog: &mut Catalog) {
let mv_qnames: BTreeSet<QualifiedName> = catalog
.materialized_views
.iter()
.map(|mv| mv.qname.clone())
.collect();
if mv_qnames.is_empty() {
return;
}
for idx in &mut catalog.indexes {
if let IndexParent::Table(qname) = &idx.on
&& mv_qnames.contains(qname)
{
idx.on = IndexParent::Mv(qname.clone());
}
}
}
struct KnownObjects {
relations: BTreeSet<QualifiedName>,
}
impl KnownObjects {
fn from_catalog(catalog: &Catalog) -> Self {
let mut relations = BTreeSet::new();
for t in &catalog.tables {
relations.insert(t.qname.clone());
}
for v in &catalog.views {
relations.insert(v.qname.clone());
}
for m in &catalog.materialized_views {
relations.insert(m.qname.clone());
}
Self { relations }
}
fn has_relation(&self, qname: &QualifiedName) -> bool {
self.relations.contains(qname)
}
}
fn walk_body_ast(
raw_body: &str,
view_qname: &QualifiedName,
known: &KnownObjects,
) -> Result<(Vec<DepEdge>, Vec<ViewColumn>), AstCanonError> {
let qname_str = view_qname.to_string();
let parsed = pg_query::parse(raw_body).map_err(|e| AstCanonError::NormalizeFailed {
view: qname_str.clone(),
reason: e.to_string(),
})?;
let mut deps: Vec<DepEdge> = Vec::new();
let mut columns: Vec<ViewColumn> = Vec::new();
let Some(root) = parsed.protobuf.stmts.first() else {
return Ok((deps, columns));
};
let Some(node) = &root.stmt else {
return Ok((deps, columns));
};
walk_node(
node,
view_qname,
known,
&mut deps,
&mut columns,
true,
&qname_str,
)?;
deps.sort();
deps.dedup();
Ok((deps, columns))
}
fn walk_node(
node: &pg_query::protobuf::Node,
view_qname: &QualifiedName,
known: &KnownObjects,
deps: &mut Vec<DepEdge>,
columns: &mut Vec<ViewColumn>,
is_top_level_select: bool,
qname_str: &str,
) -> Result<(), AstCanonError> {
use pg_query::NodeEnum as N;
let Some(inner) = &node.node else {
return Ok(());
};
match inner {
N::SelectStmt(sel) => {
if is_top_level_select {
for target in &sel.target_list {
if let Some(col) = derive_column_name(target) {
columns.push(col);
}
}
}
for from in &sel.from_clause {
walk_node(from, view_qname, known, deps, columns, false, qname_str)?;
}
if let Some(wc) = &sel.where_clause {
walk_node(wc, view_qname, known, deps, columns, false, qname_str)?;
}
if let Some(larg) = &sel.larg {
walk_select(larg, view_qname, known, deps, qname_str)?;
}
if let Some(rarg) = &sel.rarg {
walk_select(rarg, view_qname, known, deps, qname_str)?;
}
if let Some(with) = &sel.with_clause {
for cte in &with.ctes {
walk_node(cte, view_qname, known, deps, columns, false, qname_str)?;
}
}
}
N::RangeVar(rv) => {
if !rv.schemaname.is_empty()
&& !rv.relname.is_empty()
&& let Ok(s) = Identifier::from_unquoted(&rv.schemaname)
.or_else(|_| Identifier::from_quoted(&rv.schemaname))
&& let Ok(n) = Identifier::from_unquoted(&rv.relname)
.or_else(|_| Identifier::from_quoted(&rv.relname))
{
let ref_qname = QualifiedName::new(s, n);
if !known.has_relation(&ref_qname) {
return Err(AstCanonError::UnresolvedReference {
view: qname_str.to_string(),
object: ref_qname.to_string(),
});
}
deps.push(DepEdge {
from: NodeId::Table(view_qname.clone()),
to: NodeId::Table(ref_qname),
source: DepSource::AstExtracted,
});
}
}
N::JoinExpr(j) => {
if let Some(larg) = &j.larg {
walk_node(larg, view_qname, known, deps, columns, false, qname_str)?;
}
if let Some(rarg) = &j.rarg {
walk_node(rarg, view_qname, known, deps, columns, false, qname_str)?;
}
}
N::RangeSubselect(sub) => {
if let Some(subquery) = &sub.subquery {
walk_node(subquery, view_qname, known, deps, columns, false, qname_str)?;
}
}
N::CommonTableExpr(cte) => {
if let Some(q) = &cte.ctequery {
walk_node(q, view_qname, known, deps, columns, false, qname_str)?;
}
}
_ => {}
}
Ok(())
}
fn walk_select(
sel: &pg_query::protobuf::SelectStmt,
view_qname: &QualifiedName,
known: &KnownObjects,
deps: &mut Vec<DepEdge>,
qname_str: &str,
) -> Result<(), AstCanonError> {
let node = pg_query::protobuf::Node {
node: Some(pg_query::NodeEnum::SelectStmt(Box::new(sel.clone()))),
};
walk_node(
&node,
view_qname,
known,
deps,
&mut Vec::new(),
false,
qname_str,
)
}
fn resolve_explicit_column_types(columns: &mut [ViewColumn]) {
for col in columns {
if col.column_type.is_none() {
col.column_type = Some(expression_placeholder());
}
}
}
fn expression_placeholder() -> ColumnType {
ColumnType::Other {
raw: "expression".to_string(),
}
}
fn derive_column_name(target: &pg_query::protobuf::Node) -> Option<ViewColumn> {
use pg_query::NodeEnum as N;
let Some(N::ResTarget(rt)) = &target.node else {
return None;
};
if !rt.name.is_empty() {
let name = Identifier::from_unquoted(&rt.name)
.or_else(|_| Identifier::from_quoted(&rt.name))
.ok()?;
return Some(ViewColumn {
name,
column_type: Some(expression_placeholder()),
comment: None,
});
}
if let Some(val_box) = &rt.val
&& let Some(col_name) = extract_column_ref_name(val_box)
{
return Some(ViewColumn {
name: col_name,
column_type: Some(expression_placeholder()),
comment: None,
});
}
Identifier::from_quoted("?column?")
.ok()
.map(|name| ViewColumn {
name,
column_type: Some(expression_placeholder()),
comment: None,
})
}
fn extract_column_ref_name(node: &pg_query::protobuf::Node) -> Option<Identifier> {
use pg_query::NodeEnum as N;
match &node.node {
Some(N::ColumnRef(cr)) => {
for field in cr.fields.iter().rev() {
if let Some(N::String(s)) = &field.node
&& !s.sval.is_empty()
{
return Identifier::from_unquoted(&s.sval)
.or_else(|_| Identifier::from_quoted(&s.sval))
.ok();
}
}
None
}
Some(N::TypeCast(tc)) => tc.arg.as_deref().and_then(extract_column_ref_name),
_ => None,
}
}
fn apply_view_column_aliases(body_sql: &str, aliases: &[String]) -> Option<String> {
use pg_query::NodeEnum;
let parsed = pg_query::parse(body_sql).ok()?;
let mut result = parsed.protobuf;
for stmt in &mut result.stmts {
let Some(node) = stmt.stmt.as_mut().and_then(|n| n.node.as_mut()) else {
continue;
};
if let NodeEnum::SelectStmt(sel) = node {
apply_aliases_to_select(sel, aliases)?;
return pg_query::deparse(&result).ok();
}
}
None
}
fn apply_aliases_to_select(
sel: &mut pg_query::protobuf::SelectStmt,
aliases: &[String],
) -> Option<()> {
use pg_query::NodeEnum;
if sel.target_list.len() != aliases.len() {
return None;
}
for (node, alias) in sel.target_list.iter_mut().zip(aliases.iter()) {
let Some(NodeEnum::ResTarget(rt)) = node.node.as_mut() else {
return None;
};
let implicit = restarget_implicit_column_name(rt);
if implicit.as_deref() == Some(alias.as_str()) {
rt.name.clear();
} else {
rt.name.clone_from(alias);
}
}
Some(())
}
fn restarget_implicit_column_name(rt: &pg_query::protobuf::ResTarget) -> Option<String> {
use pg_query::NodeEnum;
let val = rt.val.as_ref()?;
let node = val.node.as_ref()?;
let NodeEnum::ColumnRef(cref) = node else {
return None;
};
let last = cref.fields.last()?;
let last_node = last.node.as_ref()?;
if let NodeEnum::String(s) = last_node {
Some(s.sval.clone())
} else {
None
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ir::catalog::Catalog;
use crate::ir::index::{
Index, IndexColumn, IndexColumnExpr, IndexMethod, NullsOrder, SortOrder,
};
use crate::ir::view::MaterializedView;
use crate::parse::normalize_body::NormalizedBody;
fn id(s: &str) -> Identifier {
Identifier::from_unquoted(s).unwrap()
}
fn qn(schema: &str, name: &str) -> QualifiedName {
QualifiedName::new(id(schema), id(name))
}
fn simple_index(name: &str, on: IndexParent, unique: bool) -> Index {
Index {
qname: qn("app", name),
on,
method: IndexMethod::BTree,
columns: vec![IndexColumn {
expr: IndexColumnExpr::Column(id("k")),
collation: None,
opclass: None,
sort_order: SortOrder::Asc,
nulls_order: NullsOrder::NullsLast,
}],
include: vec![],
unique,
nulls_not_distinct: false,
predicate: None,
tablespace: None,
comment: None,
storage: crate::ir::reloptions::IndexStorageOptions::default(),
}
}
fn simple_mv(schema: &str, name: &str) -> MaterializedView {
MaterializedView {
qname: qn(schema, name),
columns: vec![],
body_canonical: NormalizedBody::from_sql("SELECT 1").unwrap(),
body_dependencies: vec![],
comment: None,
raw_body: String::new(),
owner: None,
grants: vec![],
storage: crate::ir::reloptions::MaterializedViewStorageOptions::default(),
}
}
#[test]
fn unique_index_on_mv_is_promoted_to_mv_parent() {
let mv_qname = qn("app", "mv");
let mut catalog = Catalog::empty();
catalog.materialized_views.push(simple_mv("app", "mv"));
catalog.indexes.push(simple_index(
"mv_uk",
IndexParent::Table(mv_qname.clone()),
true,
));
promote_mv_index_parents(&mut catalog);
assert!(
matches!(&catalog.indexes[0].on, IndexParent::Mv(q) if q == &mv_qname),
"expected IndexParent::Mv after promotion, got {:?}",
catalog.indexes[0].on,
);
}
#[test]
fn index_on_table_is_not_promoted() {
let table_qname = qn("app", "users");
let mut catalog = Catalog::empty();
catalog.indexes.push(simple_index(
"users_idx",
IndexParent::Table(table_qname.clone()),
false,
));
promote_mv_index_parents(&mut catalog);
assert!(
matches!(&catalog.indexes[0].on, IndexParent::Table(q) if q == &table_qname),
"Table index must not be promoted: {:?}",
catalog.indexes[0].on,
);
}
#[test]
fn no_mvs_no_promotion() {
let mut catalog = Catalog::empty();
catalog.indexes.push(simple_index(
"some_idx",
IndexParent::Table(qn("app", "t")),
false,
));
promote_mv_index_parents(&mut catalog);
assert!(matches!(&catalog.indexes[0].on, IndexParent::Table(_)));
}
#[test]
fn already_mv_parent_is_unchanged() {
let mv_qname = qn("app", "mv");
let mut catalog = Catalog::empty();
catalog.materialized_views.push(simple_mv("app", "mv"));
catalog.indexes.push(simple_index(
"mv_idx",
IndexParent::Mv(mv_qname.clone()),
false,
));
promote_mv_index_parents(&mut catalog);
assert!(
matches!(&catalog.indexes[0].on, IndexParent::Mv(q) if q == &mv_qname),
"already-Mv parent must remain Mv: {:?}",
catalog.indexes[0].on,
);
}
}