pgmt 0.5.0

use crate::catalog::file_dependencies::FileDependencyAugmentation;
use crate::catalog::id::{DbObjectId, DependsOn};
use crate::diff::operations::MigrationStep;
use crate::diff::{
    aggregates as aggregates_diff, casts as casts_diff, constraints as constraints_diff,
    custom_types as custom_types_diff, domains as domains_diff, functions as functions_diff,
    indexes as indexes_diff, operators as operators_diff, policies as policies_diff,
    sequences as sequences_diff, tables as tables_diff, triggers as triggers_diff,
    views as views_diff,
};
use sqlx::PgPool;
use std::collections::BTreeMap;

pub mod aggregate;
pub mod attached;
pub mod cast;
pub mod constraint;
pub mod custom_type;
pub mod domain;
pub mod extension;
pub mod file_dependencies;
pub mod function;
pub mod grant;
pub mod id;
pub mod identity;
pub mod index;
pub mod operator;
pub mod policy;
pub mod schema;
pub mod sequence;
pub mod table;
pub mod target;
pub mod triggers;
pub mod utils;
pub mod view;

#[derive(Debug, Clone)]
pub struct Catalog {
    pub schemas: Vec<schema::Schema>,
    pub tables: Vec<table::Table>,
    pub views: Vec<view::View>,
    pub types: Vec<custom_type::CustomType>,
    pub domains: Vec<domain::Domain>,
    pub functions: Vec<function::Function>,
    pub aggregates: Vec<aggregate::Aggregate>,
    pub operators: Vec<operator::Operator>,
    pub casts: Vec<cast::Cast>,
    pub sequences: Vec<sequence::Sequence>,
    pub indexes: Vec<index::Index>,
    pub constraints: Vec<constraint::Constraint>,
    pub triggers: Vec<triggers::Trigger>,
    pub policies: Vec<policy::Policy>,
    pub extensions: Vec<extension::Extension>,
    pub grants: Vec<grant::Grant>,

    pub forward_deps: BTreeMap<DbObjectId, Vec<DbObjectId>>,
    pub reverse_deps: BTreeMap<DbObjectId, Vec<DbObjectId>>,
}

impl Catalog {
    /// Load only the objects pgmt manages: the physical catalog scoped by the
    /// objects config. This is what diff/render/validate consumers should
    /// use — comparing or rendering an unfiltered catalog leaks image-provided
    /// substrate (the shadow branch legitimately contains it).
    pub async fn load_managed(
        pool: &PgPool,
        filter: &crate::config::filter::ObjectFilter,
    ) -> anyhow::Result<Self> {
        Ok(filter.filter_catalog(Self::load_unfiltered(pool).await?))
    }

    /// Load the raw physical catalog, including objects outside pgmt's managed
    /// universe. Only for callers that genuinely need the physical world
    /// (substrate detection, init import before scoping is decided) or that
    /// provably filter downstream.
    pub async fn load_unfiltered(pool: &PgPool) -> anyhow::Result<Self> {
        Self::load_with_file_dependencies(pool, None).await
    }

    /// Load catalog with optional file-based dependency augmentation
    #[allow(clippy::explicit_auto_deref)] // Required for PoolConnection -> PgConnection deref
    pub async fn load_with_file_dependencies(
        pool: &PgPool,
        file_augmentation: Option<&FileDependencyAugmentation>,
    ) -> anyhow::Result<Self> {
        // Acquire a single connection to ensure consistent search_path across all fetches.
        // This is critical because pg_get_function_identity_arguments() output depends on
        // the connection's search_path, and we need functions and grants to match.
        let mut conn = pool.acquire().await?;

        // Set consistent search_path for all queries on this connection
        sqlx::query("SET search_path = public, pg_catalog")
            .execute(&mut *conn)
            .await?;

        let schemas = schema::fetch(&mut *conn).await?;
        let tables = table::fetch(&mut *conn).await?;
        let views = view::fetch(&mut *conn).await?;
        let types = custom_type::fetch(&mut *conn).await?;
        let domains = domain::fetch(&mut *conn).await?;
        let functions = function::fetch(&mut *conn).await?;
        let aggregates = aggregate::fetch(&mut *conn).await?;
        let operators = operator::fetch(&mut *conn).await?;
        let casts = cast::fetch(&mut *conn).await?;
        let sequences = sequence::fetch(&mut *conn).await?;
        let indexes = index::fetch(&mut *conn).await?;
        let constraints = constraint::fetch(&mut *conn).await?;
        let triggers = triggers::fetch(&mut *conn).await?;
        let policies = policy::fetch(&mut *conn).await?;
        let extensions = extension::fetch(&mut *conn).await?;
        let grants = grant::fetch(&mut *conn).await?;

        let mut forward = BTreeMap::new();
        let mut reverse = BTreeMap::new();

        fn insert_deps<T: DependsOn>(
            items: &[T],
            fwd: &mut BTreeMap<DbObjectId, Vec<DbObjectId>>,
            rev: &mut BTreeMap<DbObjectId, Vec<DbObjectId>>,
        ) {
            for item in items {
                let id = item.id();
                let deps = item.depends_on();
                fwd.insert(id.clone(), deps.to_vec());

                for dep in deps {
                    rev.entry(dep.clone()).or_default().push(id.clone());
                }
            }
        }

        insert_deps(&tables, &mut forward, &mut reverse);
        insert_deps(&views, &mut forward, &mut reverse);
        insert_deps(&types, &mut forward, &mut reverse);
        insert_deps(&domains, &mut forward, &mut reverse);
        insert_deps(&functions, &mut forward, &mut reverse);
        insert_deps(&aggregates, &mut forward, &mut reverse);
        insert_deps(&operators, &mut forward, &mut reverse);
        insert_deps(&casts, &mut forward, &mut reverse);
        insert_deps(&sequences, &mut forward, &mut reverse);
        insert_deps(&indexes, &mut forward, &mut reverse);
        insert_deps(&constraints, &mut forward, &mut reverse);
        insert_deps(&triggers, &mut forward, &mut reverse);
        insert_deps(&policies, &mut forward, &mut reverse);
        insert_deps(&extensions, &mut forward, &mut reverse);
        insert_deps(&grants, &mut forward, &mut reverse);

        let mut catalog = Self {
            schemas,
            tables,
            views,
            types,
            domains,
            functions,
            aggregates,
            operators,
            casts,
            sequences,
            indexes,
            constraints,
            triggers,
            policies,
            extensions,
            grants,
            forward_deps: forward,
            reverse_deps: reverse,
        };

        // Insert function-based casts as intermediates between their function and
        // anything that depends on that function (closes the view->cast ordering
        // gap that pg_depend can't express). Run before file augmentation so
        // user-declared `-- require:` edges layer on top.
        catalog.apply_cast_function_routing();

        if let Some(augmentation) = file_augmentation {
            catalog.apply_file_augmentation(augmentation);
        }

        Ok(catalog)
    }

    /// Create a new catalog with file-based dependencies augmented
    pub fn with_file_dependencies_augmented(
        mut self,
        augmentation: FileDependencyAugmentation,
    ) -> Self {
        self.apply_file_augmentation(&augmentation);
        self
    }

    fn apply_file_augmentation(&mut self, augmentation: &FileDependencyAugmentation) {
        for (object_id, additional_deps) in &augmentation.additional_dependencies {
            let existing_deps = self.forward_deps.entry(object_id.clone()).or_default();

            for additional_dep in additional_deps {
                if !existing_deps.contains(additional_dep) {
                    existing_deps.push(additional_dep.clone());
                }
            }
        }

        self.rebuild_reverse_deps();
    }

    /// Rebuild `reverse_deps` from the current `forward_deps`.
    pub(crate) fn rebuild_reverse_deps(&mut self) {
        self.reverse_deps.clear();
        for (object_id, deps) in &self.forward_deps {
            for dep in deps {
                self.reverse_deps
                    .entry(dep.clone())
                    .or_default()
                    .push(object_id.clone());
            }
        }
    }

    /// Route function-based casts so anything that depends on a cast's
    /// implementing function also depends on the cast itself.
    ///
    /// PostgreSQL records no `view -> cast` edge in `pg_depend`: a query that
    /// applies a function-based cast records a dependency on the cast's
    /// *function*, not on the `pg_cast` entry (see `diff::casts` and the cast
    /// catalog docs). But the cast must exist before such a view/function is
    /// created, or `CREATE` fails to resolve the cast. Since we *do* know the
    /// `consumer -> function` edge, we insert the cast as an intermediate:
    /// `consumer -> cast -> function`.
    ///
    /// This over-connects: a consumer that calls the function *directly* (not via
    /// the cast) also gains the edge. That is harmless for create-ordering (the
    /// cast is simply created a little early) and only costs an extra
    /// drop/recreate if the cast itself changes. We never attach a cast to
    /// another cast (or to itself), which keeps the rule cycle-free in practice;
    /// any pathological cycle (e.g. a cast's operand type whose CHECK calls the
    /// cast's function) is still caught by the ordering cycle detector.
    fn apply_cast_function_routing(&mut self) {
        // Collect (consumer -> cast) edges first; we mutate forward_deps after.
        let mut new_edges: Vec<(DbObjectId, DbObjectId)> = Vec::new();
        for cast in &self.casts {
            let cast_id = cast.id();
            for dep in &cast.depends_on {
                if !matches!(dep, DbObjectId::Function { .. }) {
                    continue;
                }
                if let Some(consumers) = self.reverse_deps.get(dep) {
                    for consumer in consumers {
                        if consumer != &cast_id && !matches!(consumer, DbObjectId::Cast { .. }) {
                            new_edges.push((consumer.clone(), cast_id.clone()));
                        }
                    }
                }
            }
        }

        if new_edges.is_empty() {
            return;
        }

        for (consumer, cast_id) in new_edges {
            let deps = self.forward_deps.entry(consumer).or_default();
            if !deps.contains(&cast_id) {
                deps.push(cast_id);
            }
        }

        self.rebuild_reverse_deps();
    }

    pub fn find_view(&self, schema: &str, name: &str) -> Option<&view::View> {
        self.views
            .iter()
            .find(|v| v.schema == schema && v.name == name)
    }

    pub fn find_table(&self, schema: &str, name: &str) -> Option<&table::Table> {
        self.tables
            .iter()
            .find(|t| t.schema == schema && t.name == name)
    }

    pub fn find_policy(&self, schema: &str, table: &str, name: &str) -> Option<&policy::Policy> {
        self.policies
            .iter()
            .find(|p| p.schema == schema && p.table_name == table && p.name == name)
    }

    pub fn find_constraint(
        &self,
        schema: &str,
        table: &str,
        name: &str,
    ) -> Option<&constraint::Constraint> {
        self.constraints
            .iter()
            .find(|c| c.schema == schema && c.table_name == table && c.name == name)
    }

    pub fn find_function(
        &self,
        schema: &str,
        name: &str,
        arguments: &str,
    ) -> Option<&function::Function> {
        self.functions
            .iter()
            .find(|f| f.schema == schema && f.name == name && f.arguments == arguments)
    }

    /// Find a function by signature (schema, name, and parameter types).
    /// Unlike `find_function`, this ignores parameter names and only matches on types.
    /// This is needed because PostgreSQL considers functions with the same parameter
    /// types but different parameter names to be the same function.
    pub fn find_function_by_signature(
        &self,
        reference: &function::Function,
    ) -> Option<&function::Function> {
        self.functions.iter().find(|f| {
            f.schema == reference.schema
                && f.name == reference.name
                && f.return_type == reference.return_type
                && f.parameters.len() == reference.parameters.len()
                && f.parameters
                    .iter()
                    .zip(reference.parameters.iter())
                    .all(|(a, b)| a.data_type == b.data_type && a.mode == b.mode)
        })
    }

    pub fn find_custom_type(&self, schema: &str, name: &str) -> Option<&custom_type::CustomType> {
        self.types
            .iter()
            .find(|t| t.schema == schema && t.name == name)
    }

    pub fn find_trigger(
        &self,
        schema: &str,
        table: &str,
        name: &str,
    ) -> Option<&triggers::Trigger> {
        self.triggers
            .iter()
            .find(|t| t.schema == schema && t.table_name == table && t.name == name)
    }

    pub fn find_index(&self, schema: &str, name: &str) -> Option<&index::Index> {
        self.indexes
            .iter()
            .find(|i| i.schema == schema && i.name == name)
    }

    pub fn find_domain(&self, schema: &str, name: &str) -> Option<&domain::Domain> {
        self.domains
            .iter()
            .find(|d| d.schema == schema && d.name == name)
    }

    pub fn find_sequence(&self, schema: &str, name: &str) -> Option<&sequence::Sequence> {
        self.sequences
            .iter()
            .find(|s| s.schema == schema && s.name == name)
    }

    pub fn find_aggregate(
        &self,
        schema: &str,
        name: &str,
        arguments: &str,
    ) -> Option<&aggregate::Aggregate> {
        self.aggregates
            .iter()
            .find(|a| a.schema == schema && a.name == name && a.arguments == arguments)
    }

    pub fn find_operator(
        &self,
        schema: &str,
        name: &str,
        arguments: &str,
    ) -> Option<&operator::Operator> {
        self.operators
            .iter()
            .find(|o| o.schema == schema && o.name == name && o.arguments == arguments)
    }

    pub fn find_cast(&self, source: &str, target: &str) -> Option<&cast::Cast> {
        self.casts
            .iter()
            .find(|c| c.source == source && c.target == target)
    }

    /// Every object that carries attached state (comments). Enumerated in ONE
    /// place via an exhaustive destructure: adding a field to `Catalog` fails to
    /// compile here until you decide whether the new object type is `Attached`.
    pub fn attached_objects(&self) -> Vec<&dyn crate::catalog::attached::Attached> {
        use crate::catalog::attached::Attached;
        let Catalog {
            schemas,
            tables,
            views,
            types,
            domains,
            functions,
            aggregates,
            operators,
            casts,
            sequences,
            indexes,
            constraints,
            triggers,
            policies,
            extensions,
            // Not object-attached comment state: grants are their own diff, and
            // the dep maps are derived. A new object field belongs above, not here.
            grants: _,
            forward_deps: _,
            reverse_deps: _,
        } = self;

        let mut out: Vec<&dyn Attached> = Vec::new();
        out.extend(schemas.iter().map(|x| x as &dyn Attached));
        out.extend(tables.iter().map(|x| x as &dyn Attached));
        out.extend(views.iter().map(|x| x as &dyn Attached));
        out.extend(types.iter().map(|x| x as &dyn Attached));
        out.extend(domains.iter().map(|x| x as &dyn Attached));
        out.extend(functions.iter().map(|x| x as &dyn Attached));
        out.extend(aggregates.iter().map(|x| x as &dyn Attached));
        out.extend(operators.iter().map(|x| x as &dyn Attached));
        out.extend(casts.iter().map(|x| x as &dyn Attached));
        out.extend(sequences.iter().map(|x| x as &dyn Attached));
        out.extend(indexes.iter().map(|x| x as &dyn Attached));
        out.extend(constraints.iter().map(|x| x as &dyn Attached));
        out.extend(triggers.iter().map(|x| x as &dyn Attached));
        out.extend(policies.iter().map(|x| x as &dyn Attached));
        out.extend(extensions.iter().map(|x| x as &dyn Attached));
        out
    }

    /// Synthesize DROP + CREATE steps for cascading a dependent object.
    ///
    /// Returns `None` if the object type doesn't support cascading or doesn't
    /// exist in both catalogs. Only the structural DROP/CREATE (and comments) are
    /// emitted here — the object's ACL is reapplied centrally for every recreated
    /// object by `crate::diff::cascade::reapply_acl_for_recreated_objects`, since
    /// a DROP discards all privileges regardless of what triggered the recreate.
    pub fn synthesize_drop_create(
        &self,
        id: &DbObjectId,
        new_catalog: &Catalog,
    ) -> Option<Vec<MigrationStep>> {
        let mut steps = Vec::new();

        match id {
            DbObjectId::View { schema, name } => {
                let old_view = self.find_view(schema, name)?;
                let new_view = new_catalog.find_view(schema, name)?;

                // Use diff functions for DROP and CREATE+COMMENT
                steps.extend(views_diff::diff(Some(old_view), None));
                steps.extend(views_diff::diff(None, Some(new_view)));
            }

            DbObjectId::Table { schema, name } => {
                let old_table = self.find_table(schema, name)?;
                let new_table = new_catalog.find_table(schema, name)?;

                steps.extend(tables_diff::diff(Some(old_table), None));
                steps.extend(tables_diff::diff(None, Some(new_table)));
            }

            DbObjectId::Policy {
                schema,
                table,
                name,
            } => {
                let old_policy = self.find_policy(schema, table, name)?;
                let new_policy = new_catalog.find_policy(schema, table, name)?;

                steps.extend(policies_diff::diff(Some(old_policy), None));
                steps.extend(policies_diff::diff(None, Some(new_policy)));
            }

            DbObjectId::Constraint {
                schema,
                table,
                name,
            } => {
                let old_constraint = self.find_constraint(schema, table, name)?;
                let new_constraint = new_catalog.find_constraint(schema, table, name)?;

                steps.extend(constraints_diff::diff(Some(old_constraint), None));
                steps.extend(constraints_diff::diff(None, Some(new_constraint)));
            }

            DbObjectId::Function {
                schema,
                name,
                arguments,
            }
            | DbObjectId::Procedure {
                schema,
                name,
                arguments,
            } => {
                let old_func = self.find_function(schema, name, arguments)?;
                // Use signature matching for new catalog lookup - parameter names may have
                // changed even though it's the same function (PostgreSQL identifies functions
                // by parameter types, not names)
                let new_func = new_catalog.find_function_by_signature(old_func)?;

                steps.extend(functions_diff::diff(Some(old_func), None));
                steps.extend(functions_diff::diff(None, Some(new_func)));
            }

            DbObjectId::Trigger {
                schema,
                table,
                name,
            } => {
                let old_trigger = self.find_trigger(schema, table, name)?;
                let new_trigger = new_catalog.find_trigger(schema, table, name)?;

                steps.extend(triggers_diff::diff(Some(old_trigger), None));
                steps.extend(triggers_diff::diff(None, Some(new_trigger)));
            }

            DbObjectId::Type { schema, name } => {
                let old_type = self.find_custom_type(schema, name)?;
                let new_type = new_catalog.find_custom_type(schema, name)?;

                steps.extend(custom_types_diff::diff(Some(old_type), None));
                steps.extend(custom_types_diff::diff(None, Some(new_type)));
            }

            DbObjectId::Domain { schema, name } => {
                let old = self.find_domain(schema, name)?;
                let new = new_catalog.find_domain(schema, name)?;
                steps.extend(domains_diff::diff(Some(old), None));
                steps.extend(domains_diff::diff(None, Some(new)));
            }

            DbObjectId::Index { schema, name } => {
                let old = self.find_index(schema, name)?;
                let new = new_catalog.find_index(schema, name)?;
                steps.extend(indexes_diff::diff(Some(old), None));
                steps.extend(indexes_diff::diff(None, Some(new)));
            }

            DbObjectId::Sequence { schema, name } => {
                let old = self.find_sequence(schema, name)?;
                let new = new_catalog.find_sequence(schema, name)?;
                steps.extend(sequences_diff::diff(Some(old), None));
                steps.extend(sequences_diff::diff(None, Some(new)));
            }

            DbObjectId::Aggregate {
                schema,
                name,
                arguments,
            } => {
                let old = self.find_aggregate(schema, name, arguments)?;
                let new = new_catalog.find_aggregate(schema, name, arguments)?;
                steps.extend(aggregates_diff::diff(Some(old), None));
                steps.extend(aggregates_diff::diff(None, Some(new)));
            }

            DbObjectId::Operator {
                schema,
                name,
                arguments,
            } => {
                let old = self.find_operator(schema, name, arguments)?;
                let new = new_catalog.find_operator(schema, name, arguments)?;
                steps.extend(operators_diff::diff(Some(old), None));
                steps.extend(operators_diff::diff(None, Some(new)));
            }

            DbObjectId::Cast { source, target } => {
                let old = self.find_cast(source, target)?;
                let new = new_catalog.find_cast(source, target)?;
                steps.extend(casts_diff::diff(Some(old), None));
                steps.extend(casts_diff::diff(None, Some(new)));
            }

            DbObjectId::Schema { .. }
            | DbObjectId::Extension { .. }
            | DbObjectId::Grant { .. }
            | DbObjectId::Comment { .. }
            | DbObjectId::Column { .. } => return None,
        }

        // No ACL here — it's re-stated centrally for every recreated object; see
        // the method doc and `cascade::reapply_acl_for_recreated_objects`.
        Some(steps)
    }

    /// Create an empty catalog for baseline generation
    pub fn empty() -> Self {
        Self {
            schemas: Vec::new(),
            tables: Vec::new(),
            views: Vec::new(),
            types: Vec::new(),
            domains: Vec::new(),
            functions: Vec::new(),
            aggregates: Vec::new(),
            operators: Vec::new(),
            casts: Vec::new(),
            sequences: Vec::new(),
            indexes: Vec::new(),
            constraints: Vec::new(),
            triggers: Vec::new(),
            policies: Vec::new(),
            extensions: Vec::new(),
            grants: Vec::new(),
            forward_deps: BTreeMap::new(),
            reverse_deps: BTreeMap::new(),
        }
    }

    /// Check if the catalog contains an object with the given ID
    pub fn contains_id(&self, id: &DbObjectId) -> bool {
        match id {
            DbObjectId::Schema { name } => self.schemas.iter().any(|s| &s.name == name),
            DbObjectId::Table { schema, name } => self.find_table(schema, name).is_some(),
            DbObjectId::View { schema, name } => self.find_view(schema, name).is_some(),
            DbObjectId::Type { schema, name } => self.find_custom_type(schema, name).is_some(),
            DbObjectId::Domain { schema, name } => self.find_domain(schema, name).is_some(),
            DbObjectId::Function {
                schema,
                name,
                arguments,
            }
            | DbObjectId::Procedure {
                schema,
                name,
                arguments,
            } => self.find_function(schema, name, arguments).is_some(),
            DbObjectId::Aggregate {
                schema,
                name,
                arguments,
            } => self.find_aggregate(schema, name, arguments).is_some(),
            DbObjectId::Operator {
                schema,
                name,
                arguments,
            } => self.find_operator(schema, name, arguments).is_some(),
            DbObjectId::Cast { source, target } => self.find_cast(source, target).is_some(),
            DbObjectId::Sequence { schema, name } => self.find_sequence(schema, name).is_some(),
            DbObjectId::Index { schema, name } => self.find_index(schema, name).is_some(),
            DbObjectId::Constraint {
                schema,
                table,
                name,
            } => self.find_constraint(schema, table, name).is_some(),
            DbObjectId::Trigger {
                schema,
                table,
                name,
            } => self.find_trigger(schema, table, name).is_some(),
            DbObjectId::Policy {
                schema,
                table,
                name,
            } => self.find_policy(schema, table, name).is_some(),
            DbObjectId::Extension { name } => self.extensions.iter().any(|e| &e.name == name),
            DbObjectId::Grant { id } => self.grants.iter().any(|g| &g.id() == id),
            DbObjectId::Comment { object_id } => self.contains_id(object_id),
            // Column resolves to its parent table for containment checks
            DbObjectId::Column { schema, table, .. } => self.find_table(schema, table).is_some(),
        }
    }
}