linter: fix adding-not-nullable-field for pg >= 12 with validated CHECK

crates/squawk_linter/src/rules/adding_not_null_field.rs

@@ -1,28 +1,166 @@
+use std::collections::{HashMap, HashSet};
+
 use squawk_syntax::{
-    Parse, SourceFile,
+    Parse, SourceFile, SyntaxKind,
     ast::{self, AstNode},
+    identifier::Identifier,
 };
 
-use crate::{Linter, Rule, Violation};
+use crate::{Linter, Rule, Version, Violation};
+
+/// A table/column pair identifying a NOT NULL check constraint.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+struct TableColumn {
+    table: Identifier,
+    column: Identifier,
+}
+
+/// Checks if a Literal expression is NULL.
+fn is_null_literal(lit: &ast::Literal) -> bool {
+    lit.syntax()
+        .first_child_or_token()
+        .is_some_and(|child| child.kind() == SyntaxKind::NULL_KW)
+}
+
+/// Checks if a `BinExpr` has an IS NOT operator.
+fn has_is_not_operator(bin_expr: &ast::BinExpr) -> bool {
+    // First try the Op accessor (in case the parser wraps the operator)
+    if let Some(op) = bin_expr.op() {
+        if op.is_not().is_some() {
+            return true;
+        }
+    }
+    // Fall back to checking for IS_NOT directly as a child
+    bin_expr
+        .syntax()
+        .children()
+        .any(|child| child.kind() == SyntaxKind::IS_NOT)
+}
+
+/// Checks if a CHECK constraint expression is of the form `column IS NOT NULL`.
+/// Returns the column name if so.
+fn is_not_null_check(expr: &ast::Expr) -> Option<Identifier> {
+    let ast::Expr::BinExpr(bin_expr) = expr else {
+        return None;
+    };
+    // Check if this is an `IS NOT` operation
+    if !has_is_not_operator(bin_expr) {
+        return None;
+    }
+    // Check if the RHS is NULL
+    let rhs = bin_expr.rhs()?;
+    if !matches!(rhs, ast::Expr::Literal(ref lit) if is_null_literal(lit)) {
+        return None;
+    }
+    // Get the column name from the LHS
+    let lhs = bin_expr.lhs()?;
+    match lhs {
+        ast::Expr::NameRef(name_ref) => Some(Identifier::new(&name_ref.text())),
+        _ => None,
+    }
+}
+
+/// Extracts the table name from an ALTER TABLE statement.
+fn get_table_name(alter_table: &ast::AlterTable) -> Option<Identifier> {
+    alter_table
+        .relation_name()?
+        .path()?
+        .segment()?
+        .name_ref()
+        .map(|x| Identifier::new(&x.text()))
+}
 
 pub(crate) fn adding_not_null_field(ctx: &mut Linter, parse: &Parse<SourceFile>) {
     let file = parse.tree();
+
+    // For PostgreSQL 12+, track CHECK constraints that ensure a column IS NOT NULL
+    let pg12_plus = ctx.settings.pg_version >= Version::new(12, None, None);
+
+    // Map from constraint name to the table/column it covers
+    let mut not_null_constraints: HashMap<Identifier, TableColumn> = HashMap::new();
+    // Set of table/column pairs with validated NOT NULL check constraints
+    let mut validated_not_null_columns: HashSet<TableColumn> = HashSet::new();
+
     for stmt in file.stmts() {
         if let ast::Stmt::AlterTable(alter_table) = stmt {
+            let Some(table_name) = get_table_name(&alter_table) else {
+                continue;
+            };
+
             for action in alter_table.actions() {
-                if let ast::AlterTableAction::AlterColumn(alter_column) = action {
-                    let Some(option) = alter_column.option() else {
-                        continue;
-                    };
-
-                    if matches!(option, ast::AlterColumnOption::SetNotNull(_)) {
-                        ctx.report(Violation::for_node(
-                            Rule::AddingNotNullableField,
-                            "Setting a column `NOT NULL` blocks reads while the table is scanned."
-                                .into(),
-                            option.syntax(),
-                        ).help("Make the field nullable and use a `CHECK` constraint instead."));
+                match action {
+                    // Track CHECK constraints with NOT VALID that check column IS NOT NULL
+                    ast::AlterTableAction::AddConstraint(add_constraint) if pg12_plus => {
+                        if add_constraint.not_valid().is_some() {
+                            if let Some(ast::Constraint::CheckConstraint(check)) =
+                                add_constraint.constraint()
+                            {
+                                if let Some(constraint_name) =
+                                    check.constraint_name().and_then(|c| c.name())
+                                {
+                                    if let Some(expr) = check.expr() {
+                                        if let Some(column) = is_not_null_check(&expr) {
+                                            not_null_constraints.insert(
+                                                Identifier::new(&constraint_name.text()),
+                                                TableColumn {
+                                                    table: table_name.clone(),
+                                                    column,
+                                                },
+                                            );
+                                        }
+                                    }
+                                }
+                            }
+                        }
                     }
+                    // Track VALIDATE CONSTRAINT for our tracked constraints
+                    ast::AlterTableAction::ValidateConstraint(validate_constraint) if pg12_plus => {
+                        if let Some(constraint_name) = validate_constraint
+                            .name_ref()
+                            .map(|x| Identifier::new(&x.text()))
+                        {
+                            if let Some(table_column) = not_null_constraints.get(&constraint_name) {
+                                // Only mark as validated if it's for the same table
+                                if table_column.table == table_name {
+                                    validated_not_null_columns.insert(table_column.clone());
+                                }
+                            }
+                        }
+                    }
+                    // Check SET NOT NULL
+                    ast::AlterTableAction::AlterColumn(alter_column) => {
+                        let Some(option) = alter_column.option() else {
+                            continue;
+                        };
+
+                        if matches!(option, ast::AlterColumnOption::SetNotNull(_)) {
+                            // For PostgreSQL 12+, skip if there's a validated CHECK constraint
+                            if pg12_plus {
+                                if let Some(column_name) =
+                                    alter_column.name_ref().map(|x| Identifier::new(&x.text()))
+                                {
+                                    let table_column = TableColumn {
+                                        table: table_name.clone(),
+                                        column: column_name,
+                                    };
+                                    if validated_not_null_columns.contains(&table_column) {
+                                        continue;
+                                    }
+                                }
+                            }
+
+                            ctx.report(
+                                Violation::for_node(
+                                    Rule::AddingNotNullableField,
+                                    "Setting a column `NOT NULL` blocks reads while the table is scanned."
+                                        .into(),
+                                    option.syntax(),
+                                )
+                                .help("Make the field nullable and use a `CHECK` constraint instead."),
+                            );
+                        }
+                    }
+                    _ => {}
                 }
             }
         }
@@ -33,8 +171,8 @@ pub(crate) fn adding_not_null_field(ctx: &mut Linter, parse: &Parse<SourceFile>)
 mod test {
     use insta::assert_snapshot;
 
-    use crate::Rule;
-    use crate::test_utils::{lint_errors, lint_ok};
+    use crate::test_utils::{lint_errors, lint_errors_with, lint_ok, lint_ok_with};
+    use crate::{LinterSettings, Rule};
 
     #[test]
     fn set_not_null() {
@@ -90,4 +228,137 @@ COMMIT;
         "#;
         assert_snapshot!(lint_errors(sql, Rule::AddingNotNullableField));
     }
+
+    // GitHub issue #628: SET NOT NULL should be safe on PostgreSQL 12+ when
+    // there's a validated CHECK constraint for the column.
+    #[test]
+    fn regression_gh_issue_628_pg16_with_validated_check_ok() {
+        let sql = r#"
+BEGIN;
+
+ALTER TABLE foo ADD COLUMN bar BIGINT;
+
+ALTER TABLE foo ADD CONSTRAINT bar_not_null CHECK (bar IS NOT NULL) NOT VALID;
+
+COMMIT;
+
+BEGIN;
+
+ALTER TABLE foo VALIDATE CONSTRAINT bar_not_null;
+
+ALTER TABLE foo ALTER COLUMN bar SET NOT NULL;
+
+ALTER TABLE foo DROP CONSTRAINT bar_not_null;
+
+COMMIT;
+        "#;
+
+        lint_ok_with(
+            sql,
+            LinterSettings {
+                pg_version: "16".parse().expect("Invalid PostgreSQL version"),
+                ..Default::default()
+            },
+            Rule::AddingNotNullableField,
+        );
+    }
+
+    #[test]
+    fn regression_gh_issue_628_pg12_with_validated_check_ok() {
+        let sql = r#"
+BEGIN;
+ALTER TABLE foo ADD COLUMN bar BIGINT;
+ALTER TABLE foo ADD CONSTRAINT bar_not_null CHECK (bar IS NOT NULL) NOT VALID;
+COMMIT;
+
+BEGIN;
+ALTER TABLE foo VALIDATE CONSTRAINT bar_not_null;
+ALTER TABLE foo ALTER COLUMN bar SET NOT NULL;
+ALTER TABLE foo DROP CONSTRAINT bar_not_null;
+COMMIT;
+        "#;
+        lint_ok_with(
+            sql,
+            LinterSettings {
+                pg_version: "12".parse().expect("Invalid PostgreSQL version"),
+                ..Default::default()
+            },
+            Rule::AddingNotNullableField,
+        );
+    }
+
+    #[test]
+    fn regression_gh_issue_628_pg11_with_validated_check_err() {
+        // PostgreSQL 11 doesn't support using CHECK constraint to skip table scan
+        let sql = r#"
+BEGIN;
+ALTER TABLE foo ADD COLUMN bar BIGINT;
+ALTER TABLE foo ADD CONSTRAINT bar_not_null CHECK (bar IS NOT NULL) NOT VALID;
+COMMIT;
+
+BEGIN;
+ALTER TABLE foo VALIDATE CONSTRAINT bar_not_null;
+ALTER TABLE foo ALTER COLUMN bar SET NOT NULL;
+COMMIT;
+        "#;
+        assert_snapshot!(lint_errors_with(
+            sql,
+            LinterSettings {
+                pg_version: "11".parse().expect("Invalid PostgreSQL version"),
+                ..Default::default()
+            },
+            Rule::AddingNotNullableField
+        ));
+    }
+
+    #[test]
+    fn pg12_without_validated_check_err() {
+        // Without a validated CHECK constraint, SET NOT NULL is still unsafe
+        let sql = r#"
+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 pg12_with_check_but_not_validated_err() {
+        // CHECK constraint exists but not validated yet
+        let sql = r#"
+ALTER TABLE foo ADD CONSTRAINT bar_not_null CHECK (bar IS NOT NULL) NOT VALID;
+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 pg12_with_different_column_validated_err() {
+        // Validated CHECK exists for a different column
+        let sql = r#"
+ALTER TABLE foo ADD CONSTRAINT baz_not_null CHECK (baz IS NOT NULL) NOT VALID;
+ALTER TABLE foo VALIDATE CONSTRAINT baz_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
+        ));
+    }
 }

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

@@ -0,0 +1,11 @@
+---
+source: crates/squawk_linter/src/rules/adding_not_null_field.rs
+assertion_line: 338
+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_with_different_column_validated_err.snap

@@ -0,0 +1,11 @@
+---
+source: crates/squawk_linter/src/rules/adding_not_null_field.rs
+assertion_line: 356
+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.
+  ╭▸ 
+4 │ 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_without_validated_check_err.snap

@@ -0,0 +1,11 @@
+---
+source: crates/squawk_linter/src/rules/adding_not_null_field.rs
+assertion_line: 321
+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.
+  ╭▸ 
+2 │ 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__regression_gh_issue_628_pg11_with_validated_check_err.snap

@@ -0,0 +1,11 @@
+---
+source: crates/squawk_linter/src/rules/adding_not_null_field.rs
+assertion_line: 305
+expression: "lint_errors_with(sql, LinterSettings\n{\n    pg_version: \"11\".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.
+  ╭▸ 
+9 │ ALTER TABLE foo ALTER COLUMN bar SET NOT NULL;
+  │                                  ━━━━━━━━━━━━
+  │
+  ╰ help: Make the field nullable and use a `CHECK` constraint instead.