linter: scope cross-migration SET NOT NULL suppression

crates/squawk_linter/src/rules/adding_not_null_field.rs

@@ -1,7 +1,7 @@
 use rustc_hash::{FxHashMap, FxHashSet};
 
 use squawk_syntax::{
-    Parse, SourceFile,
+    Parse, SourceFile, SyntaxNode,
     ast::{self, AstNode},
 };
 
@@ -13,6 +13,100 @@ struct TableColumn {
     column: String,
 }
 
+#[derive(Default)]
+struct NotNullValidation {
+    not_null_constraints: FxHashMap<String, TableColumn>,
+    validated_columns: FxHashSet<TableColumn>,
+    external_validates: FxHashSet<(String, String)>,
+    dropped_constraints: FxHashSet<(String, String)>,
+}
+
+impl NotNullValidation {
+    fn record_not_null_constraint(&mut self, name: String, table_column: TableColumn) {
+        self.not_null_constraints.insert(name, table_column);
+    }
+
+    fn record_validate(&mut self, table: String, constraint: String) {
+        if let Some(table_column) = self.not_null_constraints.get(&constraint)
+            && table_column.table == table
+        {
+            self.validated_columns.insert(table_column.clone());
+        } else {
+            self.external_validates.insert((table, constraint));
+        }
+    }
+
+    fn record_drop(&mut self, table: String, constraint: String) {
+        self.dropped_constraints.insert((table, constraint));
+    }
+
+    fn is_column_validated(&self, table_column: &TableColumn) -> bool {
+        self.validated_columns.contains(table_column)
+    }
+
+    fn has_external_validate_for(&self, table: &str) -> bool {
+        self.external_validates.iter().any(|(t, _)| t == table)
+    }
+
+    // External validate+drop pairs suppress SET NOT NULL warnings. When the
+    // constraint name follows the documented `<column>_not_null` /
+    // `<table>_<column>_not_null` convention we infer the exact column;
+    // otherwise we fall back to a per-table count and pair by source order.
+    fn resolved_pairs(&self) -> ResolvedPairs {
+        let mut pairs = ResolvedPairs::default();
+        for (table, constraint) in &self.external_validates {
+            if !self
+                .dropped_constraints
+                .contains(&(table.clone(), constraint.clone()))
+            {
+                continue;
+            }
+            match infer_column_from_constraint_name(table, constraint) {
+                Some(column) => {
+                    pairs.precise.insert(TableColumn {
+                        table: table.clone(),
+                        column,
+                    });
+                }
+                None => {
+                    *pairs.generic_per_table.entry(table.clone()).or_default() += 1;
+                }
+            }
+        }
+        pairs
+    }
+}
+
+#[derive(Default)]
+struct ResolvedPairs {
+    precise: FxHashSet<TableColumn>,
+    generic_per_table: FxHashMap<String, usize>,
+}
+
+// Infer the column covered by a NOT NULL helper constraint from its name.
+// Recognized forms: `<column>_not_null` and `<table>_<column>_not_null`.
+// Returns None for names that don't fit, in which case the caller falls back
+// to a generic per-table count.
+fn infer_column_from_constraint_name(table: &str, constraint: &str) -> Option<String> {
+    let stem = constraint.strip_suffix("_not_null")?;
+    if stem.is_empty() {
+        return None;
+    }
+    let table_prefix = format!("{}_", table);
+    if let Some(column) = stem.strip_prefix(&table_prefix) {
+        if column.is_empty() {
+            return None;
+        }
+        return Some(column.to_string());
+    }
+    // Bare `<stem>_not_null`: ambiguous when stem == table name (could be a
+    // table-level helper rather than a column-level one), so fall back.
+    if stem == table {
+        return None;
+    }
+    Some(stem.to_string())
+}
+
 fn is_not_null_check(expr: &ast::Expr) -> Option<String> {
     let ast::Expr::BinExpr(bin_expr) = expr else {
         return None;
@@ -48,11 +142,8 @@ pub(crate) fn adding_not_null_field(ctx: &mut Linter, parse: &Parse<SourceFile>)
 
     let is_pg12_plus = ctx.settings.pg_version >= Version::new(12, None, None);
 
-    let mut not_null_constraints: FxHashMap<String, TableColumn> = FxHashMap::default();
-    let mut validated_not_null_columns: FxHashSet<TableColumn> = FxHashSet::default();
-    // Tables where VALIDATE CONSTRAINT was seen without a matching ADD CONSTRAINT
-    // in the same file (cross-migration pattern).
-    let mut tables_with_external_validated_constraints: FxHashSet<String> = FxHashSet::default();
+    let mut validation = NotNullValidation::default();
+    let mut deferred_violations: Vec<(TableColumn, SyntaxNode)> = Vec::new();
 
     for stmt in file.stmts() {
         if let ast::Stmt::AlterTable(alter_table) = stmt {
@@ -73,7 +164,7 @@ pub(crate) fn adding_not_null_field(ctx: &mut Linter, parse: &Parse<SourceFile>)
                             && let Some(expr) = check.expr()
                             && let Some(column) = is_not_null_check(&expr)
                         {
-                            not_null_constraints.insert(
+                            validation.record_not_null_constraint(
                                 constraint_name.text(),
                                 TableColumn {
                                     table: table.clone(),
@@ -89,16 +180,15 @@ pub(crate) fn adding_not_null_field(ctx: &mut Linter, parse: &Parse<SourceFile>)
                         if let Some(constraint_name) =
                             validate_constraint.name_ref().map(|x| x.text())
                         {
-                            if let Some(table_column) = not_null_constraints.get(&constraint_name)
-                                && table_column.table == table
-                            {
-                                validated_not_null_columns.insert(table_column.clone());
-                            } else {
-                                // Cross-migration pattern: the ADD CONSTRAINT ... NOT VALID
-                                // was in a previous migration file. Track the table so that
-                                // a subsequent SET NOT NULL is considered safe.
-                                tables_with_external_validated_constraints.insert(table.clone());
-                            }
+                            validation.record_validate(table.clone(), constraint_name);
+                        }
+                    }
+                    // Track DROP CONSTRAINT for cross-migration pairing
+                    ast::AlterTableAction::DropConstraint(drop_constraint) if is_pg12_plus => {
+                        if let Some(constraint_name) =
+                            drop_constraint.name_ref().map(|x| x.text())
+                        {
+                            validation.record_drop(table.clone(), constraint_name);
                         }
                     }
                     // Step 3: Check that we're altering a validated constraint
@@ -116,9 +206,13 @@ pub(crate) fn adding_not_null_field(ctx: &mut Linter, parse: &Parse<SourceFile>)
                                 table: table.clone(),
                                 column,
                             };
-                            if validated_not_null_columns.contains(&table_column)
-                                || tables_with_external_validated_constraints.contains(&table)
-                            {
+                            if validation.is_column_validated(&table_column) {
+                                continue;
+                            }
+                            // Defer if there are external validates on this table —
+                            // we need to see the full file before deciding.
+                            if validation.has_external_validate_for(&table) {
+                                deferred_violations.push((table_column, option.syntax().clone()));
                                 continue;
                             }
                         }
@@ -138,6 +232,28 @@ pub(crate) fn adding_not_null_field(ctx: &mut Linter, parse: &Parse<SourceFile>)
             }
         }
     }
+
+    let resolved = validation.resolved_pairs();
+    let mut generic_per_table = resolved.generic_per_table;
+    for (table_column, node) in &deferred_violations {
+        if resolved.precise.contains(table_column) {
+            continue;
+        }
+        if let Some(count) = generic_per_table.get_mut(&table_column.table)
+            && *count > 0
+        {
+            *count -= 1;
+            continue;
+        }
+        ctx.report(
+            Violation::for_node(
+                Rule::AddingNotNullableField,
+                "Setting a column `NOT NULL` blocks reads while the table is scanned.".into(),
+                node,
+            )
+            .help("Make the field nullable and use a `CHECK` constraint instead."),
+        );
+    }
 }
 
 #[cfg(test)]
@@ -311,7 +427,8 @@ ALTER TABLE foo ALTER COLUMN bar SET NOT NULL;
     }
 
     // Cross-migration pattern: ADD CONSTRAINT was in a previous migration file,
-    // so only VALIDATE CONSTRAINT + SET NOT NULL appear in this file.
+    // so only VALIDATE CONSTRAINT + SET NOT NULL + DROP CONSTRAINT appear in this file.
+    // The validate+drop pairing signals it was a NOT NULL helper constraint.
     #[test]
     fn pg16_cross_migration_validate_then_set_not_null_ok() {
         let sql = r#"
@@ -351,6 +468,79 @@ ALTER TABLE foo DROP CONSTRAINT foo_bar_not_null;
         );
     }
 
+    // Validating an unrelated constraint should NOT suppress SET NOT NULL warnings.
+    #[test]
+    fn pg12_cross_migration_unrelated_validate_err() {
+        let sql = r#"
+ALTER TABLE foo VALIDATE CONSTRAINT some_other_check;
+ALTER TABLE foo ALTER COLUMN bar SET NOT NULL;
+        "#;
+        assert_snapshot!(lint_errors_with(
+            sql,
+            LinterSettings {
+                pg_version: "12".parse().expect("Invalid PostgreSQL version"),
+                ..Default::default()
+            },
+            Rule::AddingNotNullableField
+        ));
+    }
+
+    // When the constraint name follows the `<column>_not_null` convention, the
+    // pair only suppresses the matching column — other SET NOT NULL statements
+    // on the same table still warn regardless of source order.
+    #[test]
+    fn pg12_cross_migration_named_constraint_matches_column_only_err() {
+        let sql = r#"
+ALTER TABLE foo VALIDATE CONSTRAINT bar_not_null;
+ALTER TABLE foo ALTER COLUMN baz SET NOT NULL;
+ALTER TABLE foo ALTER COLUMN bar SET NOT NULL;
+ALTER TABLE foo DROP CONSTRAINT bar_not_null;
+        "#;
+        assert_snapshot!(lint_errors_with(
+            sql,
+            LinterSettings {
+                pg_version: "12".parse().expect("Invalid PostgreSQL version"),
+                ..Default::default()
+            },
+            Rule::AddingNotNullableField
+        ));
+    }
+
+    // `<table>_<column>_not_null` is also recognized by the inference.
+    #[test]
+    fn pg12_cross_migration_table_prefixed_constraint_ok() {
+        let sql = r#"
+ALTER TABLE foo VALIDATE CONSTRAINT foo_bar_not_null;
+ALTER TABLE foo ALTER COLUMN bar SET NOT NULL;
+ALTER TABLE foo DROP CONSTRAINT foo_bar_not_null;
+        "#;
+        lint_ok_with(
+            sql,
+            LinterSettings {
+                pg_version: "12".parse().expect("Invalid PostgreSQL version"),
+                ..Default::default()
+            },
+            Rule::AddingNotNullableField,
+        );
+    }
+
+    // Validate without a corresponding DROP is not the helper pattern.
+    #[test]
+    fn pg12_cross_migration_validate_no_drop_err() {
+        let sql = r#"
+ALTER TABLE foo VALIDATE CONSTRAINT bar_not_null;
+ALTER TABLE foo ALTER COLUMN bar SET NOT NULL;
+        "#;
+        assert_snapshot!(lint_errors_with(
+            sql,
+            LinterSettings {
+                pg_version: "12".parse().expect("Invalid PostgreSQL version"),
+                ..Default::default()
+            },
+            Rule::AddingNotNullableField
+        ));
+    }
+
     #[test]
     fn pg11_cross_migration_validate_then_set_not_null_err() {
         // PostgreSQL 11 doesn't support using CHECK constraint to skip table scan

crates/squawk_linter/src/rules/snapshots/squawk_linter__rules__adding_not_null_field__test__pg12_cross_migration_named_constraint_matches_column_only_err.snap

@@ -0,0 +1,10 @@
+---
+source: crates/squawk_linter/src/rules/adding_not_null_field.rs
+expression: "lint_errors_with(sql, LinterSettings\n{\n    pg_version: \"12\".parse().expect(\"Invalid PostgreSQL version\"),\n    ..Default::default()\n}, Rule::AddingNotNullableField)"
+---
+warning[adding-not-nullable-field]: Setting a column `NOT NULL` blocks reads while the table is scanned.
+  ╭▸ 
+3 │ ALTER TABLE foo ALTER COLUMN baz SET NOT NULL;
+  │                                  ━━━━━━━━━━━━
+  │
+  ╰ help: Make the field nullable and use a `CHECK` constraint instead.

crates/squawk_linter/src/rules/snapshots/squawk_linter__rules__adding_not_null_field__test__pg12_cross_migration_unrelated_validate_err.snap

@@ -0,0 +1,10 @@
+---
+source: crates/squawk_linter/src/rules/adding_not_null_field.rs
+expression: "lint_errors_with(sql, LinterSettings\n{\n    pg_version: \"12\".parse().expect(\"Invalid PostgreSQL version\"),\n    ..Default::default()\n}, Rule::AddingNotNullableField)"
+---
+warning[adding-not-nullable-field]: Setting a column `NOT NULL` blocks reads while the table is scanned.
+  ╭▸ 
+3 │ ALTER TABLE foo ALTER COLUMN bar SET NOT NULL;
+  │                                  ━━━━━━━━━━━━
+  │
+  ╰ help: Make the field nullable and use a `CHECK` constraint instead.

crates/squawk_linter/src/rules/snapshots/squawk_linter__rules__adding_not_null_field__test__pg12_cross_migration_validate_no_drop_err.snap

@@ -0,0 +1,10 @@
+---
+source: crates/squawk_linter/src/rules/adding_not_null_field.rs
+expression: "lint_errors_with(sql, LinterSettings\n{\n    pg_version: \"12\".parse().expect(\"Invalid PostgreSQL version\"),\n    ..Default::default()\n}, Rule::AddingNotNullableField)"
+---
+warning[adding-not-nullable-field]: Setting a column `NOT NULL` blocks reads while the table is scanned.
+  ╭▸ 
+3 │ ALTER TABLE foo ALTER COLUMN bar SET NOT NULL;
+  │                                  ━━━━━━━━━━━━
+  │
+  ╰ help: Make the field nullable and use a `CHECK` constraint instead.

docs/docs/adding-not-nullable-field.md

@@ -46,6 +46,14 @@ For each step, note that:
 
 See ["How not valid constraints work"](constraint-missing-not-valid.md#how-not-valid-validate-works) for more information on adding constraints as `NOT VALID`.
 
+### cross-migration suppression and constraint naming
+
+When the `ADD CONSTRAINT ... NOT VALID` lives in an earlier migration file, the rule looks for a `VALIDATE CONSTRAINT` paired with a `DROP CONSTRAINT` of the same name in the migration that runs `SET NOT NULL` to recognize the safe pattern.
+
+To get column-precise suppression (so an unrelated `SET NOT NULL` on the same table still warns), name the helper constraint `<column>_not_null` or `<table>_<column>_not_null` — for example `view_count_not_null` or `recipe_view_count_not_null`. With those names, the rule infers the exact column from the constraint name.
+
+If you use a different naming scheme, the rule falls back to a per-table count: each validate+drop pair on a table can suppress one `SET NOT NULL` on that table, paired in source order.
+
 
 ## solution for alembic and sqlalchemy