use std::collections::{BTreeSet, HashMap};
use crate::ir::catalog::Catalog;
use crate::ir::column_type::ColumnType;
use crate::parse::error::SourceLocation;
use crate::plan::edges::{DepEdge, DepSource, NodeId};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AstResolutionError {
pub message: String,
pub location: Option<SourceLocation>,
}
impl std::fmt::Display for AstResolutionError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
if let Some(loc) = &self.location {
write!(f, "{}: {}", loc, self.message)
} else {
write!(f, "{}", self.message)
}
}
}
pub fn resolve(
catalog: &Catalog,
locations: &HashMap<String, SourceLocation>,
) -> Result<(), Vec<AstResolutionError>> {
let mut errors = Vec::new();
resolve_fk_references(catalog, locations, &mut errors);
resolve_default_sequence_references(catalog, locations, &mut errors);
resolve_user_defined_references(catalog, locations, &mut errors);
resolve_routine_references(catalog, locations, &mut errors);
if errors.is_empty() {
Ok(())
} else {
Err(errors)
}
}
fn resolve_fk_references(
catalog: &Catalog,
locations: &HashMap<String, SourceLocation>,
errors: &mut Vec<AstResolutionError>,
) {
use crate::ir::constraint::ConstraintKind;
let known_tables: std::collections::BTreeSet<_> =
catalog.tables.iter().map(|t| t.qname.to_string()).collect();
for table in &catalog.tables {
for constraint in &table.constraints {
if let ConstraintKind::ForeignKey(fk) = &constraint.kind {
let ref_key = fk.referenced_table.to_string();
if !known_tables.contains(&ref_key) {
let location = locations.get(&table.qname.to_string()).cloned();
errors.push(AstResolutionError {
message: format!(
"FK {}.{} references {} which is not declared in source",
table.qname, constraint.qname.name, fk.referenced_table,
),
location,
});
}
}
}
}
}
fn resolve_default_sequence_references(
catalog: &Catalog,
locations: &HashMap<String, SourceLocation>,
errors: &mut Vec<AstResolutionError>,
) {
use crate::ir::default_expr::DefaultExpr;
let known_sequences: BTreeSet<_> = catalog
.sequences
.iter()
.map(|s| s.qname.to_string())
.collect();
for table in &catalog.tables {
for column in &table.columns {
if let Some(DefaultExpr::Sequence(qname)) = &column.default {
let seq_key = qname.to_string();
if !known_sequences.contains(&seq_key) {
let location = locations.get(&table.qname.to_string()).cloned();
errors.push(AstResolutionError {
message: format!(
"column {}.{} defaults to nextval({}) but sequence is not declared",
table.qname, column.name, qname,
),
location,
});
}
}
}
}
}
fn resolve_user_defined_references(
catalog: &Catalog,
locations: &HashMap<String, SourceLocation>,
errors: &mut Vec<AstResolutionError>,
) {
use crate::ir::user_type::UserTypeKind;
let known_types: BTreeSet<String> = catalog.types.iter().map(|t| t.qname.to_string()).collect();
for table in &catalog.tables {
for column in &table.columns {
if let ColumnType::UserDefined(qname) = &column.ty {
let key = qname.to_string();
if !known_types.contains(&key) {
let location = locations.get(&table.qname.to_string()).cloned();
errors.push(AstResolutionError {
message: format!(
"column {}.{} has type {} which is not declared in source",
table.qname, column.name, qname,
),
location,
});
}
}
}
}
for ut in &catalog.types {
match &ut.kind {
UserTypeKind::Composite { attributes } => {
for attr in attributes {
if let ColumnType::UserDefined(qname) = &attr.ty {
let key = qname.to_string();
if !known_types.contains(&key) {
let location = locations.get(&ut.qname.to_string()).cloned();
errors.push(AstResolutionError {
message: format!(
"composite type {}.{} has type {} which is not declared in source",
ut.qname, attr.name, qname,
),
location,
});
}
}
}
}
UserTypeKind::Domain { base, .. } => {
if let ColumnType::UserDefined(qname) = base {
let key = qname.to_string();
if !known_types.contains(&key) {
let location = locations.get(&ut.qname.to_string()).cloned();
errors.push(AstResolutionError {
message: format!(
"domain {} base type {} is not declared in source",
ut.qname, qname,
),
location,
});
}
}
}
UserTypeKind::Enum { .. } => {
}
UserTypeKind::Range { subtype, .. } => {
if subtype.schema.as_str() != "pg_catalog" {
let key = subtype.to_string();
if !known_types.contains(&key) {
let location = locations.get(&ut.qname.to_string()).cloned();
errors.push(AstResolutionError {
message: format!(
"range type {} subtype {} is not declared in source",
ut.qname, subtype,
),
location,
});
}
}
}
}
}
}
fn resolve_routine_references(
catalog: &Catalog,
locations: &HashMap<String, SourceLocation>,
errors: &mut Vec<AstResolutionError>,
) {
let known_tables: BTreeSet<String> =
catalog.tables.iter().map(|t| t.qname.to_string()).collect();
let known_views: BTreeSet<String> = catalog.views.iter().map(|v| v.qname.to_string()).collect();
let known_mvs: BTreeSet<String> = catalog
.materialized_views
.iter()
.map(|mv| mv.qname.to_string())
.collect();
let known_types: BTreeSet<String> = catalog.types.iter().map(|t| t.qname.to_string()).collect();
let known_functions_by_qname: BTreeSet<String> = catalog
.functions
.iter()
.map(|f| f.qname.to_string())
.collect();
let known_procedures: BTreeSet<String> = catalog
.procedures
.iter()
.map(|p| p.qname.to_string())
.collect();
for func in &catalog.functions {
let location = locations.get(&format!(
"functions.{}({})",
func.qname,
func.args
.iter()
.filter(|a| matches!(
a.mode,
crate::ir::function::ArgMode::In
| crate::ir::function::ArgMode::InOut
| crate::ir::function::ArgMode::Variadic
))
.map(|a| a.ty.render_sql())
.collect::<Vec<_>>()
.join(",")
));
check_body_deps(
&func.body_dependencies,
&func.qname.to_string(),
location,
&known_tables,
&known_views,
&known_mvs,
&known_types,
&known_functions_by_qname,
&known_procedures,
errors,
);
}
for proc in &catalog.procedures {
let location = locations.get(&format!("procedures.{}", proc.qname));
check_body_deps(
&proc.body_dependencies,
&proc.qname.to_string(),
location,
&known_tables,
&known_views,
&known_mvs,
&known_types,
&known_functions_by_qname,
&known_procedures,
errors,
);
}
}
#[allow(clippy::too_many_arguments)]
fn check_body_deps(
deps: &[DepEdge],
routine_qname: &str,
location: Option<&SourceLocation>,
known_tables: &BTreeSet<String>,
known_views: &BTreeSet<String>,
known_mvs: &BTreeSet<String>,
known_types: &BTreeSet<String>,
known_functions_by_qname: &BTreeSet<String>,
known_procedures: &BTreeSet<String>,
errors: &mut Vec<AstResolutionError>,
) {
for dep in deps {
match dep.source {
DepSource::AstExtracted => {
if let NodeId::Table(ref target_qname) = dep.to {
let key = target_qname.to_string();
if !known_tables.contains(&key)
&& !known_views.contains(&key)
&& !known_mvs.contains(&key)
{
errors.push(AstResolutionError {
message: format!(
"routine {routine_qname} body references {key} which is not \
declared in source (not a table, view, or materialized view)"
),
location: location.cloned(),
});
}
}
}
DepSource::AstDeclared => {
if let NodeId::Table(ref target_qname) = dep.to {
let key = target_qname.to_string();
let found = known_tables.contains(&key)
|| known_views.contains(&key)
|| known_mvs.contains(&key)
|| known_types.contains(&key)
|| known_functions_by_qname.contains(&key)
|| known_procedures.contains(&key);
if !found {
errors.push(AstResolutionError {
message: format!(
"routine {routine_qname} has `-- @pgevolve dep: {key}` directive \
but {key} is not declared in source (probed tables, views, \
materialized views, types, functions, and procedures)"
),
location: location.cloned(),
});
}
}
}
DepSource::Structural => {
}
}
}
}