pgmt 0.6.4

//! The logical `Table` and its columns.
//!
//! Tables are fetched and converted in [`crate::catalog::raw::table`]: the
//! catalog query keeps OIDs and attnums, and the converter resolves them into
//! the name-based identities this module's structs carry.

use super::collation::CollationRef;
use super::id::{DbObjectId, DependsOn};

#[derive(Debug, Clone)]
pub struct Column {
    pub name: String,
    pub data_type: String,
    pub default: Option<String>,
    pub not_null: bool,
    pub generated: Option<String>,
    /// `GENERATED { ALWAYS | BY DEFAULT } AS IDENTITY`. The identity-owned
    /// sequence is internal to the column (pg_depend 'i'), not a standalone
    /// catalog object.
    pub identity: Option<IdentityKind>,
    /// Explicit collation, schema-qualified. `None` when the column uses its
    /// type's default collation, so plain `text` columns never carry a
    /// `COLLATE` clause. System collations (pg_catalog."C", etc.) are stored
    /// and rendered but add no managed dependency.
    pub collation: Option<CollationRef>,
    pub comment: Option<String>,
    /// Dependencies for this column (e.g., functions used in generated expression)
    pub depends_on: Vec<DbObjectId>,
}

/// Which identity flavor a column uses (`pg_attribute.attidentity`)
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum IdentityKind {
    /// `GENERATED ALWAYS AS IDENTITY` (attidentity = 'a')
    Always,
    /// `GENERATED BY DEFAULT AS IDENTITY` (attidentity = 'd')
    ByDefault,
}

impl IdentityKind {
    pub fn as_sql(&self) -> &'static str {
        match self {
            IdentityKind::Always => "ALWAYS",
            IdentityKind::ByDefault => "BY DEFAULT",
        }
    }

    pub(crate) fn from_attidentity(attidentity: Option<&str>) -> Option<Self> {
        match attidentity {
            Some("a") => Some(IdentityKind::Always),
            Some("d") => Some(IdentityKind::ByDefault),
            _ => None,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PrimaryKey {
    pub name: String,
    pub columns: Vec<String>,
    pub comment: Option<String>,
}

#[derive(Debug, Clone)]
pub struct Table {
    pub schema: String,
    pub name: String,
    pub columns: Vec<Column>,
    pub primary_key: Option<PrimaryKey>,
    pub comment: Option<String>,

    /// Row-level security is enabled on this table
    pub rls_enabled: bool,

    /// Row-level security is forced even for table owners
    pub rls_forced: bool,

    table_dependencies: Vec<DbObjectId>,

    all_dependencies: Vec<DbObjectId>,
}

impl Table {
    pub fn id(&self) -> DbObjectId {
        DbObjectId::Table {
            schema: self.schema.clone(),
            name: self.name.clone(),
        }
    }

    /// Compute all dependencies by aggregating table + unique column dependencies
    fn compute_all_dependencies(&self) -> Vec<DbObjectId> {
        let mut all_deps = self.table_dependencies.clone();

        for column in &self.columns {
            for col_dep in &column.depends_on {
                if !all_deps.contains(col_dep) {
                    all_deps.push(col_dep.clone());
                }
            }
        }

        all_deps
    }

    /// Update the computed all_dependencies field after columns have changed
    pub fn update_all_dependencies(&mut self) {
        self.all_dependencies = self.compute_all_dependencies();
    }

    /// Create a new Table with proper dependency computation
    pub fn new(
        schema: String,
        name: String,
        columns: Vec<Column>,
        primary_key: Option<PrimaryKey>,
        comment: Option<String>,
        table_dependencies: Vec<DbObjectId>,
    ) -> Self {
        let mut table = Self {
            schema,
            name,
            columns,
            primary_key,
            comment,
            rls_enabled: false,
            rls_forced: false,
            table_dependencies: table_dependencies.clone(),
            all_dependencies: table_dependencies,
        };
        table.update_all_dependencies();
        table
    }
}

impl DependsOn for Table {
    fn id(&self) -> DbObjectId {
        self.id()
    }

    fn depends_on(&self) -> &[DbObjectId] {
        &self.all_dependencies
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::diff::operations::{ColumnAction, MigrationStep, TableOperation};
    use crate::diff::tables::diff;

    fn make_test_table(
        schema: &str,
        name: &str,
        columns: Vec<(&str, &str, bool)>,
        pk: Option<(&str, Vec<&str>)>,
    ) -> Table {
        let columns = columns
            .into_iter()
            .map(|(name, data_type, not_null)| Column {
                name: name.to_string(),
                data_type: data_type.to_string(),
                default: None,
                generated: None,
                identity: None,
                collation: None,
                comment: None,
                depends_on: vec![],
                not_null,
            })
            .collect();

        let primary_key = pk.map(|(name, cols)| PrimaryKey {
            name: name.to_string(),
            columns: cols.into_iter().map(|s| s.to_string()).collect(),
            comment: None,
        });

        Table::new(
            schema.to_string(),
            name.to_string(),
            columns,
            primary_key,
            None,
            vec![DbObjectId::Schema {
                name: schema.to_string(),
            }],
        )
    }

    #[test]
    fn test_create_table_with_primary_key() {
        let table = make_test_table(
            "public",
            "users",
            vec![
                ("id", "serial", true),
                ("name", "text", true),
                ("email", "text", false),
            ],
            Some(("users_pkey", vec!["id"])),
        );

        let steps = diff(None, Some(&table));

        assert_eq!(steps.len(), 1);
        match &steps[0] {
            MigrationStep::Table(TableOperation::Create {
                schema,
                name: table_name,
                columns,
                primary_key,
            }) => {
                assert_eq!(schema, "public");
                assert_eq!(table_name, "users");
                assert_eq!(columns.len(), 3);

                assert!(primary_key.is_some());
                let pk = primary_key.as_ref().unwrap();
                assert_eq!(pk.name, "users_pkey");
                assert_eq!(pk.columns, vec!["id"]);
            }
            _ => panic!("Expected CreateTable step"),
        }
    }

    #[test]
    fn test_create_table_without_primary_key() {
        let table = make_test_table(
            "public",
            "logs",
            vec![("timestamp", "timestamp", true), ("message", "text", false)],
            None,
        );

        let steps = diff(None, Some(&table));

        assert_eq!(steps.len(), 1);
        match &steps[0] {
            MigrationStep::Table(TableOperation::Create {
                schema,
                name: table_name,
                columns,
                primary_key,
            }) => {
                assert_eq!(schema, "public");
                assert_eq!(table_name, "logs");
                assert_eq!(columns.len(), 2);
                assert!(primary_key.is_none());
            }
            _ => panic!("Expected CreateTable step"),
        }
    }

    #[test]
    fn test_add_primary_key() {
        let old_table = make_test_table(
            "public",
            "users",
            vec![("id", "serial", true), ("name", "text", true)],
            None,
        );

        let new_table = make_test_table(
            "public",
            "users",
            vec![("id", "serial", true), ("name", "text", true)],
            Some(("users_pkey", vec!["id"])),
        );

        let steps = diff(Some(&old_table), Some(&new_table));

        assert_eq!(steps.len(), 1);
        match &steps[0] {
            MigrationStep::Table(TableOperation::Alter {
                schema,
                name,
                actions,
            }) => {
                assert_eq!(schema, "public");
                assert_eq!(name, "users");
                assert_eq!(actions.len(), 1);

                match &actions[0] {
                    ColumnAction::AddPrimaryKey { constraint } => {
                        assert_eq!(constraint.name, "users_pkey");
                        assert_eq!(constraint.columns, vec!["id"]);
                    }
                    _ => panic!("Expected AddPrimaryKey action"),
                }
            }
            _ => panic!("Expected AlterTable step"),
        }
    }

    #[test]
    fn test_drop_primary_key() {
        let old_table = make_test_table(
            "public",
            "users",
            vec![("id", "serial", true), ("name", "text", true)],
            Some(("users_pkey", vec!["id"])),
        );

        let new_table = make_test_table(
            "public",
            "users",
            vec![("id", "serial", true), ("name", "text", true)],
            None,
        );

        let steps = diff(Some(&old_table), Some(&new_table));

        assert_eq!(steps.len(), 1);
        match &steps[0] {
            MigrationStep::Table(TableOperation::Alter {
                schema,
                name,
                actions,
            }) => {
                assert_eq!(schema, "public");
                assert_eq!(name, "users");
                assert_eq!(actions.len(), 1);

                match &actions[0] {
                    ColumnAction::DropPrimaryKey { name } => {
                        assert_eq!(name, "users_pkey");
                    }
                    _ => panic!("Expected DropPrimaryKey action"),
                }
            }
            _ => panic!("Expected AlterTable step"),
        }
    }

    #[test]
    fn test_change_primary_key() {
        let old_table = make_test_table(
            "public",
            "users",
            vec![
                ("id", "serial", true),
                ("name", "text", true),
                ("email", "text", true),
            ],
            Some(("users_pkey", vec!["id"])),
        );

        let new_table = make_test_table(
            "public",
            "users",
            vec![
                ("id", "serial", true),
                ("name", "text", true),
                ("email", "text", true),
            ],
            Some(("users_email_pkey", vec!["email"])),
        );

        let steps = diff(Some(&old_table), Some(&new_table));

        assert_eq!(steps.len(), 1);
        match &steps[0] {
            MigrationStep::Table(TableOperation::Alter {
                schema,
                name,
                actions,
            }) => {
                assert_eq!(schema, "public");
                assert_eq!(name, "users");
                assert_eq!(actions.len(), 2);

                match &actions[0] {
                    ColumnAction::DropPrimaryKey { name } => {
                        assert_eq!(name, "users_pkey");
                    }
                    _ => panic!("Expected DropPrimaryKey action first"),
                }

                match &actions[1] {
                    ColumnAction::AddPrimaryKey { constraint } => {
                        assert_eq!(constraint.name, "users_email_pkey");
                        assert_eq!(constraint.columns, vec!["email"]);
                    }
                    _ => panic!("Expected AddPrimaryKey action second"),
                }
            }
            _ => panic!("Expected AlterTable step"),
        }
    }

    #[test]
    fn test_compound_primary_key() {
        let table = make_test_table(
            "public",
            "order_items",
            vec![
                ("order_id", "integer", true),
                ("product_id", "integer", true),
                ("quantity", "integer", true),
            ],
            Some(("order_items_pkey", vec!["order_id", "product_id"])),
        );

        let steps = diff(None, Some(&table));

        assert_eq!(steps.len(), 1);
        match &steps[0] {
            MigrationStep::Table(TableOperation::Create {
                schema,
                name: table_name,
                columns: _,
                primary_key,
            }) => {
                assert_eq!(schema, "public");
                assert_eq!(table_name, "order_items");

                assert!(primary_key.is_some());
                let pk = primary_key.as_ref().unwrap();
                assert_eq!(pk.name, "order_items_pkey");
                assert_eq!(pk.columns, vec!["order_id", "product_id"]);
            }
            _ => panic!("Expected CreateTable step"),
        }
    }
}