pgsafe 0.11.0

Static safety linter for PostgreSQL DDL migrations — catches unsafe schema changes before they lock or break production.
Documentation
use crate::ast::protobuf::ColumnDef;
use crate::ast::NodeEnum;

use super::Rule;
use crate::fix::{FixAnchor, FixDraft, FixDraftEdit};
use crate::RuleHit;

/// Serial pseudo-type → the equivalent real integer type the identity rewrite uses, preserving the
/// column's width (serial=int4, smallserial=int2, bigserial=int8).
const SERIAL_TO_INT: &[(&str, &str)] = &[
    ("serial", "int"),
    ("serial4", "int"),
    ("serial2", "smallint"),
    ("smallserial", "smallint"),
    ("serial8", "bigint"),
    ("bigserial", "bigint"),
];

pub struct PreferIdentity;

/// If `col`'s declared type is a serial pseudo-type, return `(replacement_int_type, type_token_location)`.
/// Serials always parse as a single unqualified name, so a schema-qualified (`names.len() != 1`) type is
/// never a serial — a cheap pre-filter, not a unique discriminator (`int4`/`int8` and user types are also
/// `len == 1`). The real test is the `SERIAL_TO_INT` string lookup below.
fn serial_replacement(col: &ColumnDef) -> Option<(&'static str, i32)> {
    let tn = col.type_name.as_ref()?;
    if tn.names.len() != 1 {
        return None;
    }
    let name = match tn.names.first()?.node.as_ref()? {
        NodeEnum::String(s) => s.sval.to_ascii_lowercase(),
        _ => return None,
    };
    let target = SERIAL_TO_INT
        .iter()
        .find(|(serial, _)| *serial == name)
        .map(|(_, int_type)| *int_type)?;
    Some((target, tn.location))
}

impl Rule for PreferIdentity {
    fn id(&self) -> &'static str {
        "prefer-identity"
    }
    // severity() defaults to Warning.
    fn check(&self, node: &NodeEnum, out: &mut Vec<RuleHit>) {
        // The autofix is offered only in CREATE TABLE. Rewriting an `ADD COLUMN serial` to an
        // identity column would swap one rewrite-hazard Error (add-column-serial) for another
        // (add-column-identity), which the `--fix` fixpoint loop's "no new error" guard rejects; in
        // an (empty/new) CREATE TABLE the swap is genuinely safe.
        let in_create = matches!(node, NodeEnum::CreateStmt(_));
        for col in super::defined_columns(node) {
            let Some((target, location)) = serial_replacement(col) else {
                continue;
            };
            let fix = in_create
                .then(|| u32::try_from(location).ok())
                .flatten()
                .map(|at| FixDraft {
                    title: "Use identity column",
                    edits: vec![FixDraftEdit {
                        anchor: FixAnchor::ReplaceTokenAt(at),
                        replacement: format!("{target} GENERATED BY DEFAULT AS IDENTITY"),
                    }],
                });
            out.push(RuleHit {
                message:
                    "A `serial` column creates an implicitly-owned sequence with ownership and \
                          permission sharp edges and is not SQL-standard."
                        .into(),
                guidance:
                    "Use an identity column — `GENERATED BY DEFAULT AS IDENTITY` — which owns \
                           its sequence cleanly and is SQL-standard."
                        .into(),
                fix,
            });
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::fix::apply;
    use crate::{lint_sql, LintOptions};

    fn findings(sql: &str) -> Vec<crate::Finding> {
        lint_sql(sql, &LintOptions::default()).unwrap()
    }
    fn hit(sql: &str) -> Option<crate::Finding> {
        findings(sql)
            .into_iter()
            .find(|f| f.rule_id == "prefer-identity")
    }

    #[test]
    fn flags_serial_in_create_table() {
        assert!(hit("CREATE TABLE t (id serial)").is_some());
    }
    #[test]
    fn flags_serial_in_add_column() {
        assert!(hit("ALTER TABLE t ADD COLUMN id serial").is_some());
    }
    #[test]
    fn ignores_plain_int_and_identity() {
        assert!(hit("CREATE TABLE t (id int)").is_none());
        assert!(hit("CREATE TABLE t (id bigint GENERATED BY DEFAULT AS IDENTITY)").is_none());
    }

    #[test]
    fn create_table_fix_preserves_width_and_clears() {
        let sql = "CREATE TABLE t (id serial)";
        let fix = hit(sql).unwrap().fix.expect("fix present in CREATE TABLE");
        assert_eq!(fix.title, "Use identity column");
        let fixed = apply(sql, &fix.edits);
        assert_eq!(
            fixed,
            "CREATE TABLE t (id int GENERATED BY DEFAULT AS IDENTITY)"
        );
        assert!(
            hit(&fixed).is_none(),
            "fixed SQL must not re-trigger prefer-identity"
        );
    }

    #[test]
    fn bigserial_maps_to_bigint() {
        let sql = "CREATE TABLE t (id bigserial)";
        let fixed = apply(sql, &hit(sql).unwrap().fix.unwrap().edits);
        assert_eq!(
            fixed,
            "CREATE TABLE t (id bigint GENERATED BY DEFAULT AS IDENTITY)"
        );
    }
    #[test]
    fn smallserial_maps_to_smallint() {
        let sql = "CREATE TABLE t (id smallserial)";
        let fixed = apply(sql, &hit(sql).unwrap().fix.unwrap().edits);
        assert_eq!(
            fixed,
            "CREATE TABLE t (id smallint GENERATED BY DEFAULT AS IDENTITY)"
        );
    }

    #[test]
    fn numeric_alias_widths_are_preserved() {
        // Each alias maps to a distinct width; assert the exact rewrite for all three.
        for (alias, int_type) in [
            ("serial4", "int"),
            ("serial2", "smallint"),
            ("serial8", "bigint"),
        ] {
            let sql = format!("CREATE TABLE t (id {alias})");
            let fixed = apply(&sql, &hit(&sql).unwrap().fix.unwrap().edits);
            assert_eq!(
                fixed,
                format!("CREATE TABLE t (id {int_type} GENERATED BY DEFAULT AS IDENTITY)"),
                "alias {alias}"
            );
        }
    }

    #[test]
    fn uppercase_serial_is_flagged() {
        // PostgreSQL down-cases unquoted identifiers, so unquoted SERIAL is the pseudo-type.
        assert!(hit("CREATE TABLE t (id SERIAL)").is_some());
        assert!(hit("CREATE TABLE t (id BIGSERIAL)").is_some());
    }

    #[test]
    fn multiple_serial_columns_all_fix_and_compose() {
        // Two serial columns → two distinct-width fixes that must splice together correctly.
        let sql = "CREATE TABLE t (a serial, b bigserial)";
        let edits: Vec<_> = findings(sql)
            .into_iter()
            .filter(|f| f.rule_id == "prefer-identity")
            .filter_map(|f| f.fix)
            .flat_map(|fx| fx.edits)
            .collect();
        assert_eq!(edits.len(), 2, "one edit per serial column");
        let fixed = apply(sql, &edits);
        assert_eq!(
            fixed,
            "CREATE TABLE t (a int GENERATED BY DEFAULT AS IDENTITY, \
             b bigint GENERATED BY DEFAULT AS IDENTITY)"
        );
        assert!(
            lint_sql(&fixed, &LintOptions::default())
                .unwrap()
                .iter()
                .all(|f| f.rule_id != "prefer-identity"),
            "composed fix clears all prefer-identity findings"
        );
    }

    #[test]
    fn add_column_fires_but_offers_no_fix() {
        // ADD COLUMN serial→identity would swap one rewrite-hazard Error for another, so no fix.
        let f = hit("ALTER TABLE t ADD COLUMN id serial").unwrap();
        assert!(f.fix.is_none());
    }

    #[test]
    fn serial_with_primary_key_fix_is_valid() {
        let sql = "CREATE TABLE t (id serial PRIMARY KEY)";
        let fixed = apply(sql, &hit(sql).unwrap().fix.unwrap().edits);
        assert_eq!(
            fixed,
            "CREATE TABLE t (id int GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY)"
        );
        // The rewritten SQL parses and no longer triggers prefer-identity (proves "valid").
        assert!(hit(&fixed).is_none());
    }
}