drizzle-migrations 0.1.16

//! `PostgreSQL` snapshot type matching drizzle-kit format.
//!
//! `PostgresSnapshot` is a type alias of the generic
//! [`crate::snapshot::Snapshot`] — the CRUD / serde IO surface lives once
//! in that module. This file supplies:
//!
//! * the [`SnapshotEntity`] impl pinning the Postgres dialect / version
//!   constants used by `Snapshot::new()`;
//! * Postgres-specific helpers (`scoped_to_tables`, `filter_serial_sequences`,
//!   `normalize_columns_for_push`, `table_names`, `schema_names`,
//!   `prepare_for_push`) attached via an `impl Snapshot<PostgresEntity>`
//!   block, which orphan rules permit because `PostgresEntity` is local.
//! * the legacy v7 type preserved for reading older snapshots.

use crate::postgres::ddl::PostgresEntity;
use crate::postgres::grammar::{extract_nextval_sequence, is_serial_expression};
use crate::snapshot::{Snapshot, SnapshotEntity};
use crate::version::POSTGRES_SNAPSHOT_VERSION;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};

impl SnapshotEntity for PostgresEntity {
    // drizzle-kit PG snapshots (and this repo's journal Dialect serde /
    // upgrade path) use "postgresql", not "postgres". Nothing validates the
    // string on load, so snapshots written with the old value still
    // deserialize fine.
    const DIALECT: &'static str = "postgresql";
    const SNAPSHOT_VERSION: &'static str = POSTGRES_SNAPSHOT_VERSION;
}

/// `PostgreSQL` schema snapshot (drizzle-kit beta v8 format).
///
/// Type alias of [`Snapshot<PostgresEntity>`]; see the generic type's docs
/// for the field set and IO surface. Postgres-specific filtering and
/// normalisation methods are attached below.
pub type PostgresSnapshot = Snapshot<PostgresEntity>;

// =============================================================================
// Postgres-specific snapshot operations
// =============================================================================
//
// Allowed by orphan rules because `PostgresEntity` is local to this crate;
// the alias `PostgresSnapshot = Snapshot<PostgresEntity>` resolves to this
// concrete instantiation.

impl Snapshot<PostgresEntity> {
    /// Return a new snapshot scoped to only the given tables.
    ///
    /// - Schema entities are kept only if referenced by a desired table.
    /// - Table-scoped entities (Column, Index, FK, PK, Unique, Check,
    ///   Policy) are kept only when their parent table is in the set.
    /// - Other global entities (Enum, Sequence, Role, View) pass through.
    ///
    /// The set contains `(schema, table_name)` pairs.
    #[must_use]
    pub fn scoped_to_tables(&self, tables: &HashSet<(String, String)>) -> Self {
        // Derive the set of schema names referenced by desired tables
        let schemas: HashSet<&str> = tables.iter().map(|(s, _)| s.as_str()).collect();

        let mut scoped = Self::new();
        for entity in &self.ddl {
            match entity {
                // Schema entities — keep only if referenced by desired tables
                PostgresEntity::Schema(s) => {
                    if schemas.contains(s.name.as_ref()) {
                        scoped.ddl.push(entity.clone());
                    }
                }
                // Table-scoped entities — keep only if table is in desired set
                PostgresEntity::Table(t) => {
                    if tables.contains(&(t.schema.to_string(), t.name.to_string())) {
                        scoped.ddl.push(entity.clone());
                    }
                }
                PostgresEntity::Column(c) => {
                    if tables.contains(&(c.schema.to_string(), c.table.to_string())) {
                        scoped.ddl.push(entity.clone());
                    }
                }
                PostgresEntity::Index(i) => {
                    if tables.contains(&(i.schema.to_string(), i.table.to_string())) {
                        scoped.ddl.push(entity.clone());
                    }
                }
                PostgresEntity::ForeignKey(f) => {
                    if tables.contains(&(f.schema.to_string(), f.table.to_string())) {
                        scoped.ddl.push(entity.clone());
                    }
                }
                PostgresEntity::PrimaryKey(p) => {
                    if tables.contains(&(p.schema.to_string(), p.table.to_string())) {
                        scoped.ddl.push(entity.clone());
                    }
                }
                PostgresEntity::UniqueConstraint(u) => {
                    if tables.contains(&(u.schema.to_string(), u.table.to_string())) {
                        scoped.ddl.push(entity.clone());
                    }
                }
                PostgresEntity::CheckConstraint(c) => {
                    if tables.contains(&(c.schema.to_string(), c.table.to_string())) {
                        scoped.ddl.push(entity.clone());
                    }
                }
                PostgresEntity::Policy(p) => {
                    if tables.contains(&(p.schema.to_string(), p.table.to_string())) {
                        scoped.ddl.push(entity.clone());
                    }
                }
                // Schema-scoped global entities — keep only if in relevant schemas
                PostgresEntity::Sequence(s) => {
                    if schemas.contains(s.schema.as_ref()) {
                        scoped.ddl.push(entity.clone());
                    }
                }
                PostgresEntity::Enum(e) => {
                    if schemas.contains(e.schema.as_ref()) {
                        scoped.ddl.push(entity.clone());
                    }
                }
                PostgresEntity::View(v) => {
                    if schemas.contains(v.schema.as_ref()) {
                        scoped.ddl.push(entity.clone());
                    }
                }
                // Truly global entities (Role, Privilege)
                _ => scoped.ddl.push(entity.clone()),
            }
        }
        scoped
    }

    /// Remove sequences that are owned by serial/bigserial columns.
    ///
    /// Serial columns auto-create sequences in `PostgreSQL`. These should not
    /// appear in snapshots used for diffing, otherwise the diff engine will
    /// try to DROP them (breaking the serial column) or CREATE duplicates.
    ///
    /// Superseded: introspection now drops auto-owned sequences at the source
    /// using `pg_depend` ownership (see `process_sequences`), which — unlike
    /// this name-pattern heuristic — cannot misclassify a hand-managed
    /// sequence that happens to be named `*_seq`. Kept for callers holding
    /// snapshots from other origins.
    pub fn filter_serial_sequences(&mut self) {
        self.filter_serial_sequences_except(&HashSet::new());
    }

    /// Like [`Self::filter_serial_sequences`], but sequences in `keep` are
    /// exempt from the name-pattern heuristic — a sequence the desired schema
    /// explicitly declares is hand-managed, no matter what it's named.
    pub fn filter_serial_sequences_except(&mut self, keep: &HashSet<(String, String)>) {
        // Collect (schema, seq_name) pairs referenced by serial column defaults
        let serial_seqs: HashSet<(String, String)> = self
            .ddl
            .iter()
            .filter_map(|e| {
                if let PostgresEntity::Column(c) = e {
                    let default = c.default.as_deref()?;
                    if is_serial_expression(default, &c.schema) {
                        let name = extract_nextval_sequence(default)?;
                        return Some((c.schema.to_string(), name));
                    }
                }
                None
            })
            .collect();

        if !serial_seqs.is_empty() {
            self.ddl.retain(|e| {
                if let PostgresEntity::Sequence(s) = e {
                    let key = (s.schema.to_string(), s.name.to_string());
                    keep.contains(&key) || !serial_seqs.contains(&key)
                } else {
                    true
                }
            });
        }
    }

    /// Normalize introspected columns for push comparison.
    ///
    /// - Converts `int4 + nextval()` back to `SERIAL` (and analogously for
    ///   int8→BIGSERIAL, int2→SMALLSERIAL) so the live snapshot matches the
    ///   desired snapshot that uses serial pseudo-types.
    /// - Strips `ordinal_position` from all columns (desired snapshots don't
    ///   have it but introspected ones do).
    pub fn normalize_columns_for_push(&mut self) {
        // Sequences still present as entities are real standalone sequences
        // (introspection already dropped serial/identity-owned ones); a
        // column defaulting to nextval() on one of them is NOT a serial
        // column and must keep its explicit default.
        let standalone_seqs = self.sequence_names();

        for entity in &mut self.ddl {
            if let PostgresEntity::Column(c) = entity {
                // Strip fields that only appear in introspection
                c.ordinal_position = None;
                // pg_catalog is the default for built-in types — clear it
                if c.type_schema.as_deref() == Some("pg_catalog") {
                    c.type_schema = None;
                }
                // Detect serial pattern: integer type + nextval() default
                if let Some(ref default) = c.default
                    && is_serial_expression(default, &c.schema)
                    && !extract_nextval_sequence(default)
                        .is_some_and(|seq| standalone_seqs.contains(&(c.schema.to_string(), seq)))
                {
                    let serial_type = match c.sql_type.as_ref() {
                        "int4" | "integer" => Some("SERIAL"),
                        "int8" | "bigint" => Some("BIGSERIAL"),
                        "int2" | "smallint" => Some("SMALLSERIAL"),
                        _ => None,
                    };
                    if let Some(st) = serial_type {
                        c.sql_type = st.to_string().into();
                        c.default = None;
                    }
                }
            }
        }
    }

    /// Extract the set of `(schema, name)` pairs for sequence entities in
    /// this snapshot.
    #[must_use]
    pub fn sequence_names(&self) -> HashSet<(String, String)> {
        self.ddl
            .iter()
            .filter_map(|e| {
                if let PostgresEntity::Sequence(s) = e {
                    Some((s.schema.to_string(), s.name.to_string()))
                } else {
                    None
                }
            })
            .collect()
    }

    /// Keep only sequence entities whose `(schema, name)` is in `managed`.
    ///
    /// Push uses this to leave hand-managed sequences alone: a live sequence
    /// the desired schema doesn't declare is unmanaged — dropping it would be
    /// destructive, and it must not produce a diff at all.
    pub fn retain_sequences(&mut self, managed: &HashSet<(String, String)>) {
        self.ddl.retain(|e| {
            if let PostgresEntity::Sequence(s) = e {
                managed.contains(&(s.schema.to_string(), s.name.to_string()))
            } else {
                true
            }
        });
    }

    /// Extract the set of `(schema, table_name)` pairs in this snapshot.
    #[must_use]
    pub fn table_names(&self) -> HashSet<(String, String)> {
        let mut tables = HashSet::new();
        for entity in &self.ddl {
            if let PostgresEntity::Table(t) = entity {
                tables.insert((t.schema.to_string(), t.name.to_string()));
            }
        }
        tables
    }

    /// Extract the unique schema names referenced by tables in this snapshot.
    #[must_use]
    pub fn schema_names(&self) -> Vec<String> {
        let mut names: Vec<String> = self
            .table_names()
            .into_iter()
            .map(|(s, _)| s)
            .collect::<HashSet<_>>()
            .into_iter()
            .collect();
        names.sort();
        names
    }

    /// Prepare a live (introspected) snapshot for push comparison against `desired`.
    ///
    /// Combines four normalization steps into a single call:
    /// 1. Scope to only tables present in `desired` (avoids DROP for unmanaged tables)
    /// 2. Drop serial-pattern sequences NOT declared by `desired`. Live
    ///    introspection already excludes serial/identity-owned sequences via
    ///    `pg_depend`, but snapshots from other origins may still carry them;
    ///    a sequence `desired` declares is hand-managed and must survive so
    ///    the serial detection below leaves its column's default alone.
    /// 3. Normalize columns (int4+nextval→SERIAL, strip `ordinal_position`) —
    ///    a nextval() default on a sequence that survives step 2 (a real
    ///    standalone sequence) is not mistaken for a serial column
    /// 4. Scope sequences to the ones `desired` declares (unmanaged
    ///    standalone sequences must be left alone, not dropped)
    #[must_use]
    pub fn prepare_for_push(&self, desired: &Self) -> Self {
        let tables = desired.table_names();
        let managed = desired.sequence_names();
        let mut scoped = self.scoped_to_tables(&tables);
        scoped.filter_serial_sequences_except(&managed);
        scoped.normalize_columns_for_push();
        scoped.retain_sequences(&managed);
        scoped
    }
}

// =============================================================================
// Legacy V7 Snapshot (drizzle-kit stable)
// =============================================================================

/// Schema metadata for tracking renames (Legacy)
#[derive(Serialize, Deserialize, Clone, Debug, Default)]
pub struct Meta {
    #[serde(default)]
    pub schemas: HashMap<String, String>,
    #[serde(default)]
    pub tables: HashMap<String, String>,
    #[serde(default)]
    pub columns: HashMap<String, String>,
}

/// Legacy V7 Snapshot for reading compatibility
#[derive(Serialize, Deserialize, Clone, Debug)]
#[serde(rename_all = "camelCase")]
pub struct SnapshotV7 {
    pub version: String,
    pub dialect: String,
    pub id: String,
    pub prev_id: String,
    // Using Value for legacy details to avoid redefining all legacy structs
    pub tables: HashMap<String, serde_json::Value>,
    pub enums: HashMap<String, serde_json::Value>,
    pub schemas: HashMap<String, serde_json::Value>,
    pub sequences: HashMap<String, serde_json::Value>,
    #[serde(default)]
    pub views: HashMap<String, serde_json::Value>,
    #[serde(rename = "_meta")]
    pub meta: Meta,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::postgres::ddl::{Column, Schema, Sequence, Table};
    use crate::version::ORIGIN_UUID;

    fn make_table(schema: &str, name: &str) -> PostgresEntity {
        PostgresEntity::Table(Table {
            schema: schema.to_string().into(),
            name: name.to_string().into(),
            is_unlogged: None,
            is_temporary: None,
            inherits: None,
            tablespace: None,
            is_rls_enabled: None,
            comment: None,
        })
    }

    fn make_column(schema: &str, table: &str, name: &str, sql_type: &str) -> Column {
        Column::new(
            schema.to_string(),
            table.to_string(),
            name.to_string(),
            sql_type.to_string(),
        )
    }

    fn make_sequence(schema: &str, name: &str) -> PostgresEntity {
        PostgresEntity::Sequence(Sequence {
            schema: schema.to_string().into(),
            name: name.to_string().into(),
            increment_by: None,
            min_value: None,
            max_value: None,
            start_with: None,
            cache_size: None,
            cycle: None,
        })
    }

    #[test]
    fn test_new_snapshot() {
        let snapshot = PostgresSnapshot::new();
        assert_eq!(snapshot.version, "8");
        assert_eq!(snapshot.dialect, "postgresql");
        assert_eq!(snapshot.prev_ids, vec![ORIGIN_UUID]);
        assert!(snapshot.ddl.is_empty());
        assert!(snapshot.renames.is_empty());
    }

    #[test]
    fn test_add_entity() {
        let mut snapshot = PostgresSnapshot::new();

        let schema = Schema::new("public");
        snapshot.add_entity(PostgresEntity::Schema(schema));

        let table = Table {
            schema: "public".into(),
            name: "users".into(),
            is_unlogged: None,
            is_temporary: None,
            inherits: None,
            tablespace: None,
            is_rls_enabled: None,
            comment: None,
        };
        snapshot.add_entity(PostgresEntity::Table(table));

        assert_eq!(snapshot.ddl.len(), 2);
    }

    #[test]
    fn test_schema_names() {
        let mut snap = PostgresSnapshot::new();
        snap.add_entity(make_table("public", "users"));
        snap.add_entity(make_table("auth", "sessions"));
        snap.add_entity(make_table("public", "posts"));

        let names = snap.schema_names();
        assert_eq!(names, vec!["auth", "public"]);
    }

    #[test]
    fn test_schema_names_empty() {
        let snap = PostgresSnapshot::new();
        assert!(snap.schema_names().is_empty());
    }

    #[test]
    fn test_filter_serial_sequences() {
        let mut snap = PostgresSnapshot::new();
        snap.add_entity(make_table("public", "users"));
        // Serial column with nextval default
        let mut col = make_column("public", "users", "id", "int4");
        col.default = Some("nextval('users_id_seq'::regclass)".into());
        snap.add_entity(PostgresEntity::Column(col));
        // The auto-created sequence
        snap.add_entity(make_sequence("public", "users_id_seq"));
        // An unrelated sequence that should survive
        snap.add_entity(make_sequence("public", "custom_seq"));

        snap.filter_serial_sequences();

        let seq_names: Vec<&str> = snap
            .ddl
            .iter()
            .filter_map(|e| {
                if let PostgresEntity::Sequence(s) = e {
                    Some(s.name.as_ref())
                } else {
                    None
                }
            })
            .collect();
        assert_eq!(seq_names, vec!["custom_seq"]);
    }

    #[test]
    fn test_normalize_columns_for_push() {
        let mut snap = PostgresSnapshot::new();
        // int4 + nextval → should become SERIAL
        let mut col = make_column("public", "users", "id", "int4");
        col.default = Some("nextval('users_id_seq'::regclass)".into());
        col.ordinal_position = Some(1);
        col.type_schema = Some("pg_catalog".into());
        snap.add_entity(PostgresEntity::Column(col));

        // bigint + nextval → BIGSERIAL
        let mut col2 = make_column("public", "users", "big_id", "bigint");
        col2.default = Some("nextval('users_big_id_seq'::regclass)".into());
        snap.add_entity(PostgresEntity::Column(col2));

        // Regular column — should be untouched (except ordinal/type_schema)
        let mut col3 = make_column("public", "users", "name", "text");
        col3.ordinal_position = Some(3);
        snap.add_entity(PostgresEntity::Column(col3));

        snap.normalize_columns_for_push();

        let columns: Vec<&Column> = snap
            .ddl
            .iter()
            .filter_map(|e| {
                if let PostgresEntity::Column(c) = e {
                    Some(c)
                } else {
                    None
                }
            })
            .collect();

        // id: int4+nextval → SERIAL, no default, no ordinal, no type_schema
        assert_eq!(columns[0].sql_type.as_ref(), "SERIAL");
        assert!(columns[0].default.is_none());
        assert!(columns[0].ordinal_position.is_none());
        assert!(columns[0].type_schema.is_none());

        // big_id: bigint+nextval → BIGSERIAL
        assert_eq!(columns[1].sql_type.as_ref(), "BIGSERIAL");
        assert!(columns[1].default.is_none());

        // name: unchanged type, ordinal stripped
        assert_eq!(columns[2].sql_type.as_ref(), "text");
        assert!(columns[2].ordinal_position.is_none());
    }

    #[test]
    fn test_prepare_for_push() {
        // Live snapshot has extra tables/sequences
        let mut live = PostgresSnapshot::new();
        live.add_entity(PostgresEntity::Schema(Schema::new("public")));
        live.add_entity(make_table("public", "users"));
        live.add_entity(make_table("public", "unmanaged"));
        let mut col = make_column("public", "users", "id", "int4");
        col.default = Some("nextval('users_id_seq'::regclass)".into());
        col.ordinal_position = Some(1);
        col.type_schema = Some("pg_catalog".into());
        live.add_entity(PostgresEntity::Column(col));
        live.add_entity(make_sequence("public", "users_id_seq"));

        // Desired only has "users"
        let mut desired = PostgresSnapshot::new();
        desired.add_entity(make_table("public", "users"));

        let result = live.prepare_for_push(&desired);

        // "unmanaged" table should be filtered out
        let table_names: Vec<&str> = result
            .ddl
            .iter()
            .filter_map(|e| {
                if let PostgresEntity::Table(t) = e {
                    Some(t.name.as_ref())
                } else {
                    None
                }
            })
            .collect();
        assert_eq!(table_names, vec!["users"]);

        // Serial sequence should be filtered
        let seq_count = result
            .ddl
            .iter()
            .filter(|e| matches!(e, PostgresEntity::Sequence(_)))
            .count();
        assert_eq!(seq_count, 0);

        // Column should be normalized to SERIAL
        let col = result
            .ddl
            .iter()
            .find_map(|e| {
                if let PostgresEntity::Column(c) = e {
                    Some(c)
                } else {
                    None
                }
            })
            .unwrap();
        assert_eq!(col.sql_type.as_ref(), "SERIAL");
        assert!(col.default.is_none());
        assert!(col.ordinal_position.is_none());
        assert!(col.type_schema.is_none());
    }

    #[test]
    fn test_prepare_for_push_keeps_declared_hand_managed_sequence() {
        // A standalone sequence with a serial-shaped name, used as an
        // explicit column default. Because `desired` declares it, the
        // name-pattern heuristic must not eat it and the column must NOT be
        // rewritten to SERIAL.
        let mut live = PostgresSnapshot::new();
        live.add_entity(PostgresEntity::Schema(Schema::new("public")));
        live.add_entity(make_table("public", "invoices"));
        let mut col = make_column("public", "invoices", "number", "int8");
        col.default = Some("nextval('invoices_number_seq'::regclass)".into());
        live.add_entity(PostgresEntity::Column(col));
        live.add_entity(make_sequence("public", "invoices_number_seq"));

        let mut desired = PostgresSnapshot::new();
        desired.add_entity(make_table("public", "invoices"));
        desired.add_entity(make_sequence("public", "invoices_number_seq"));

        let result = live.prepare_for_push(&desired);

        let seq_count = result
            .ddl
            .iter()
            .filter(|e| matches!(e, PostgresEntity::Sequence(_)))
            .count();
        assert_eq!(seq_count, 1, "declared hand-managed sequence must survive");

        let col = result
            .ddl
            .iter()
            .find_map(|e| {
                if let PostgresEntity::Column(c) = e {
                    Some(c)
                } else {
                    None
                }
            })
            .unwrap();
        assert_eq!(
            col.sql_type.as_ref(),
            "int8",
            "column on a hand-managed sequence must not become BIGSERIAL"
        );
        assert!(
            col.default.is_some(),
            "explicit nextval default must remain"
        );
    }

    #[test]
    fn test_prepare_for_push_drops_unmanaged_standalone_sequence() {
        // A live standalone sequence the desired schema does not declare is
        // unmanaged: it must vanish from the comparison entirely (no DROP
        // SEQUENCE diff), and since introspection guarantees it is not
        // serial-owned, its referencing column keeps its default only if the
        // sequence was declared — here it is not, so serial detection applies.
        let mut live = PostgresSnapshot::new();
        live.add_entity(PostgresEntity::Schema(Schema::new("public")));
        live.add_entity(make_table("public", "users"));
        live.add_entity(make_sequence("public", "audit_seq"));

        let mut desired = PostgresSnapshot::new();
        desired.add_entity(make_table("public", "users"));

        let result = live.prepare_for_push(&desired);

        let seq_count = result
            .ddl
            .iter()
            .filter(|e| matches!(e, PostgresEntity::Sequence(_)))
            .count();
        assert_eq!(
            seq_count, 0,
            "unmanaged sequence must not enter the diff at all"
        );
    }

    #[test]
    fn test_scoped_to_tables_keeps_relevant_entities() {
        let mut snap = PostgresSnapshot::new();
        snap.add_entity(PostgresEntity::Schema(Schema::new("public")));
        snap.add_entity(PostgresEntity::Schema(Schema::new("other")));
        snap.add_entity(make_table("public", "users"));
        snap.add_entity(make_table("other", "logs"));
        snap.add_entity(make_sequence("public", "my_seq"));
        snap.add_entity(make_sequence("other", "other_seq"));

        let tables: HashSet<(String, String)> =
            [("public".to_string(), "users".to_string())].into();
        let scoped = snap.scoped_to_tables(&tables);

        // Only "public" schema, "users" table, and "public" sequence
        let schemas: Vec<&str> = scoped
            .ddl
            .iter()
            .filter_map(|e| {
                if let PostgresEntity::Schema(s) = e {
                    Some(s.name.as_ref())
                } else {
                    None
                }
            })
            .collect();
        assert_eq!(schemas, vec!["public"]);

        let tables: Vec<&str> = scoped
            .ddl
            .iter()
            .filter_map(|e| {
                if let PostgresEntity::Table(t) = e {
                    Some(t.name.as_ref())
                } else {
                    None
                }
            })
            .collect();
        assert_eq!(tables, vec!["users"]);

        let seqs: Vec<&str> = scoped
            .ddl
            .iter()
            .filter_map(|e| {
                if let PostgresEntity::Sequence(s) = e {
                    Some(s.name.as_ref())
                } else {
                    None
                }
            })
            .collect();
        assert_eq!(seqs, vec!["my_seq"]);
    }
}