uqa_sql/schema/

constraint_metadata.rs

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Normalize durable constraint names and identities through a caller-owned identity allocator.
use std::collections::BTreeSet;
use uqa_core::RelationIdentity;

#[derive(Debug, thiserror::Error)]
#[error("{0}")]
pub struct ConstraintMetadataError(pub String);
pub type ConstraintMetadataResult<T> = Result<T, ConstraintMetadataError>;
pub type CatalogIdentityAllocator<'a> = dyn FnMut(&str) -> ConstraintMetadataResult<[u8; 16]> + 'a;

pub fn materialize_constraint_metadata(
    relation: &RelationIdentity,
    columns: &mut [crate::ast::ColumnDef],
    constraints: &mut crate::ast::TableConstraintSet,
    allocate: &mut CatalogIdentityAllocator<'_>,
) -> ConstraintMetadataResult<bool> {
    // Releases predating typed table-key persistence stored column-level PRIMARY KEY and UNIQUE declarations only as ColumnDef flags. Promote those legacy flags before assigning names so catalog publication always sees named constraints.
    let mut changed = promote_legacy_column_key_constraints(columns, constraints);
    let mut used = BTreeSet::new();
    for column in columns.iter() {
        record_constraint_name(&mut used, column.not_null_name.as_deref())?;
        record_constraint_name(&mut used, column.check_name.as_deref())?;
        record_constraint_name(
            &mut used,
            column
                .references
                .as_ref()
                .and_then(|reference| reference.name.as_deref()),
        )?;
    }
    for constraint in &constraints.key_constraints {
        record_constraint_name(&mut used, constraint.name.as_deref())?;
    }
    for constraint in &constraints.checks {
        record_constraint_name(&mut used, constraint.name.as_deref())?;
    }
    for constraint in &constraints.foreign_keys {
        record_constraint_name(&mut used, constraint.name.as_deref())?;
    }

    let mut column_object_ids = BTreeSet::new();
    for column in columns.iter_mut() {
        if column
            .object_id
            .is_some_and(|object_id| !column_object_ids.insert(object_id))
        {
            column.object_id = None;
        }
        changed |= assign_catalog_object_id(&mut column.object_id, "column", allocate)?;
        if let Some(object_id) = column.object_id {
            column_object_ids.insert(object_id);
        }
        if column.not_null {
            changed |= assign_constraint_name(
                &mut column.not_null_name,
                format!("{}_{}_not_null", relation.name, column.name),
                &mut used,
            )?;
        }
        if column.check.is_some() {
            changed |= assign_constraint_name(
                &mut column.check_name,
                format!("{}_{}_check", relation.name, column.name),
                &mut used,
            )?;
            changed |= assign_catalog_object_id(
                &mut column.check_object_id,
                "CHECK constraint",
                allocate,
            )?;
        }
        if let Some(reference) = &mut column.references {
            changed |= assign_constraint_name(
                &mut reference.name,
                format!("{}_{}_fkey", relation.name, column.name),
                &mut used,
            )?;
            changed |= assign_constraint_object_id(&mut reference.object_id, allocate)?;
        }
    }
    for constraint in &mut constraints.key_constraints {
        let base = match constraint.kind {
            crate::ast::TableKeyConstraintKind::PrimaryKey => {
                format!("{}_pkey", relation.name)
            }
            crate::ast::TableKeyConstraintKind::Unique => format!(
                "{}_{}_key",
                relation.name,
                constraint_column_component(&constraint.columns, relation)?
            ),
        };
        changed |= assign_constraint_name(&mut constraint.name, base, &mut used)?;
    }
    for constraint in &mut constraints.checks {
        let mut referenced_columns = Vec::new();
        collect_constraint_columns(&constraint.expr, &mut referenced_columns);
        let base = if referenced_columns.len() == 1 {
            format!("{}_{}_check", relation.name, referenced_columns[0])
        } else {
            format!("{}_check", relation.name)
        };
        changed |= assign_constraint_name(&mut constraint.name, base, &mut used)?;
        changed |=
            assign_catalog_object_id(&mut constraint.object_id, "CHECK constraint", allocate)?;
    }
    changed |= synchronize_partition_inherited_foreign_key_ids(constraints);
    for constraint in &mut constraints.foreign_keys {
        let component = constraint_column_component(&constraint.local_columns, relation)?;
        changed |= assign_constraint_name(
            &mut constraint.name,
            format!("{}_{}_fkey", relation.name, component),
            &mut used,
        )?;
        changed |= assign_constraint_object_id(&mut constraint.object_id, allocate)?;
    }
    changed |= synchronize_partition_inherited_foreign_key_ids(constraints);
    Ok(changed)
}

fn promote_legacy_column_key_constraints(
    columns: &[crate::ast::ColumnDef],
    constraints: &mut crate::ast::TableConstraintSet,
) -> bool {
    let mut changed = false;
    for column in columns {
        for (present, kind) in [
            (
                column.primary_key,
                crate::ast::TableKeyConstraintKind::PrimaryKey,
            ),
            (column.unique, crate::ast::TableKeyConstraintKind::Unique),
        ] {
            if !present
                || constraints.key_constraints.iter().any(|constraint| {
                    constraint.kind == kind
                        && constraint.columns.as_slice() == [column.name.as_str()]
                })
            {
                continue;
            }
            constraints
                .key_constraints
                .push(crate::ast::TableKeyConstraint {
                    name: None,
                    kind,
                    columns: vec![column.name.clone()],
                    nulls_not_distinct: false,
                    without_overlaps: false,
                });
            changed = true;
        }
    }
    changed
}

pub fn foreign_keys_match_without_object_id(
    left: &crate::ast::ForeignKey,
    right: &crate::ast::ForeignKey,
) -> bool {
    let mut left = left.clone();
    let mut right = right.clone();
    left.object_id = None;
    right.object_id = None;
    left == right
}

pub fn synchronize_partition_inherited_foreign_key_ids(
    constraints: &mut crate::ast::TableConstraintSet,
) -> bool {
    let mut changed = false;
    for inherited_index in 0..constraints.hierarchy.partition_inherited_foreign_keys.len() {
        let inherited = &constraints.hierarchy.partition_inherited_foreign_keys[inherited_index];
        let Some(foreign_key_index) = constraints
            .foreign_keys
            .iter()
            .position(|foreign_key| foreign_keys_match_without_object_id(foreign_key, inherited))
        else {
            continue;
        };
        let object_id = constraints.foreign_keys[foreign_key_index]
            .object_id
            .or(inherited.object_id);
        if constraints.foreign_keys[foreign_key_index].object_id != object_id {
            constraints.foreign_keys[foreign_key_index].object_id = object_id;
            changed = true;
        }
        if constraints.hierarchy.partition_inherited_foreign_keys[inherited_index].object_id
            != object_id
        {
            constraints.hierarchy.partition_inherited_foreign_keys[inherited_index].object_id =
                object_id;
            changed = true;
        }
    }
    changed
}

fn assign_constraint_object_id(
    target: &mut Option<[u8; 16]>,
    allocate: &mut CatalogIdentityAllocator<'_>,
) -> ConstraintMetadataResult<bool> {
    assign_catalog_object_id(target, "foreign-key constraint", allocate)
}

fn assign_catalog_object_id(
    target: &mut Option<[u8; 16]>,
    object_kind: &str,
    allocate: &mut CatalogIdentityAllocator<'_>,
) -> ConstraintMetadataResult<bool> {
    if target.is_some() {
        return Ok(false);
    }
    *target = Some(allocate(object_kind)?);
    Ok(true)
}

fn record_constraint_name(
    used: &mut BTreeSet<String>,
    name: Option<&str>,
) -> ConstraintMetadataResult<()> {
    let Some(name) = name else {
        return Ok(());
    };
    if name.is_empty() {
        return Err(ConstraintMetadataError(
            "constraint name must not be empty".into(),
        ));
    }
    if !used.insert(name.to_string()) {
        return Err(ConstraintMetadataError(format!(
            "constraint `{name}` is declared more than once"
        )));
    }
    Ok(())
}

fn assign_constraint_name(
    target: &mut Option<String>,
    base: String,
    used: &mut BTreeSet<String>,
) -> ConstraintMetadataResult<bool> {
    if target.is_some() {
        return Ok(false);
    }
    if used.insert(base.clone()) {
        *target = Some(base);
        return Ok(true);
    }
    for suffix in 1_u64.. {
        let candidate = format!("{base}{suffix}");
        if used.insert(candidate.clone()) {
            *target = Some(candidate);
            return Ok(true);
        }
    }
    Err(ConstraintMetadataError(format!(
        "constraint name suffix space exhausted for `{base}`"
    )))
}

fn constraint_column_component(
    columns: &[String],
    relation: &RelationIdentity,
) -> ConstraintMetadataResult<String> {
    if columns.is_empty() {
        return Err(ConstraintMetadataError(format!(
            "constraint on table `{}` has no columns",
            relation.qualified_name()
        )));
    }
    Ok(columns.join("_"))
}

fn collect_constraint_columns(expression: &crate::ast::Expr, output: &mut Vec<String>) {
    use crate::ast::{Expr, FrameBound};
    match expression {
        Expr::Column(name) | Expr::QualifiedColumn { column: name, .. } => {
            if !output.contains(name) {
                output.push(name.clone());
            }
        }
        Expr::Func {
            args,
            order_by,
            filter,
            ..
        } => {
            for argument in args {
                collect_constraint_columns(argument, output);
            }
            for order in order_by {
                collect_constraint_columns(&order.expr, output);
            }
            if let Some(filter) = filter {
                collect_constraint_columns(filter, output);
            }
        }
        Expr::Array(items) | Expr::Row(items) | Expr::And(items) | Expr::Or(items) => {
            for item in items {
                collect_constraint_columns(item, output);
            }
        }
        Expr::Binary { lhs, rhs, .. } => {
            collect_constraint_columns(lhs, output);
            collect_constraint_columns(rhs, output);
        }
        Expr::Not(inner)
        | Expr::UnaryMinus(inner)
        | Expr::IsNull { expr: inner, .. }
        | Expr::Cast { expr: inner, .. } => {
            collect_constraint_columns(inner, output);
        }
        Expr::Between { expr, low, high } => {
            collect_constraint_columns(expr, output);
            collect_constraint_columns(low, output);
            collect_constraint_columns(high, output);
        }
        Expr::InList { expr, list, .. } => {
            collect_constraint_columns(expr, output);
            for item in list {
                collect_constraint_columns(item, output);
            }
        }
        Expr::WindowCall { args, spec, .. } => {
            for argument in args {
                collect_constraint_columns(argument, output);
            }
            for expression in &spec.partition_by {
                collect_constraint_columns(expression, output);
            }
            for order in &spec.order_by {
                collect_constraint_columns(&order.expr, output);
            }
            if let Some(frame) = &spec.frame {
                for bound in [&frame.start, &frame.end] {
                    if let FrameBound::Preceding(expression) | FrameBound::Following(expression) =
                        bound
                    {
                        collect_constraint_columns(expression, output);
                    }
                }
            }
        }
        Expr::Case {
            base,
            when,
            else_branch,
        } => {
            if let Some(base) = base {
                collect_constraint_columns(base, output);
            }
            for (condition, result) in when {
                collect_constraint_columns(condition, output);
                collect_constraint_columns(result, output);
            }
            if let Some(else_branch) = else_branch {
                collect_constraint_columns(else_branch, output);
            }
        }
        Expr::InSubquery { expr, .. } => collect_constraint_columns(expr, output),
        Expr::Default
        | Expr::Star
        | Expr::QualifiedStar(_)
        | Expr::InternalColumn(_)
        | Expr::Literal(_)
        | Expr::TypedLiteral { .. }
        | Expr::Param(_)
        | Expr::ScalarSubquery(_)
        | Expr::Exists { .. } => {}
    }
}