From 65e53e037beb5c06dc63d6cb5145fa421281de7c Mon Sep 17 00:00:00 2001
From: Christopher Dignam <chris@dignam.xyz>
Date: Thu, 11 Jan 2024 23:59:38 -0500
Subject: [PATCH] improve unique constraint check

fixes #336
---
 .../src/rules/disallow_unique_constraint.rs   | 30 ++++++++++++++++++-
 ...__unique_constraint_inline_add_column.snap | 23 ++++++++++++++
 ...e_constraint_inline_add_column_unique.snap | 23 ++++++++++++++
 3 files changed, 75 insertions(+), 1 deletion(-)
 create mode 100644 linter/src/rules/snapshots/squawk_linter__rules__disallow_unique_constraint__test_rules__unique_constraint_inline_add_column.snap
 create mode 100644 linter/src/rules/snapshots/squawk_linter__rules__disallow_unique_constraint__test_rules__unique_constraint_inline_add_column_unique.snap

diff --git a/linter/src/rules/disallow_unique_constraint.rs b/linter/src/rules/disallow_unique_constraint.rs
index bf8fe92c..16400639 100644
--- a/linter/src/rules/disallow_unique_constraint.rs
+++ b/linter/src/rules/disallow_unique_constraint.rs
@@ -3,7 +3,7 @@ use crate::versions::Version;
 use crate::violations::{RuleViolation, RuleViolationKind};
 
 use squawk_parser::ast::{
-    AlterTableCmds, AlterTableDef, AlterTableType, ConstrType, RawStmt, Stmt,
+    AlterTableCmds, AlterTableDef, AlterTableType, ColumnDefConstraint, ConstrType, RawStmt, Stmt,
 };
 
 #[must_use]
@@ -36,6 +36,20 @@ pub fn disallow_unique_constraint(
                                 ));
                             }
                         }
+                        (Some(AlterTableDef::ColumnDef(col)), AlterTableType::AddColumn) => {
+                            if tables_created.contains(tbl_name) {
+                                continue;
+                            }
+                            for ColumnDefConstraint::Constraint(constraint) in &col.constraints {
+                                if constraint.contype == ConstrType::Unique {
+                                    errs.push(RuleViolation::new(
+                                        RuleViolationKind::DisallowedUniqueConstraint,
+                                        raw_stmt.into(),
+                                        None,
+                                    ));
+                                }
+                            }
+                        }
                         _ => continue,
                     }
                 }
@@ -144,4 +158,18 @@ ALTER TABLE products ADD CONSTRAINT sku_constraint UNIQUE (sku);
         "#;
         assert_eq!(lint_sql_assuming_in_transaction(sql), vec![]);
     }
+    #[test]
+    fn test_unique_constraint_inline_add_column() {
+        let sql = r#"
+ALTER TABLE foo ADD COLUMN bar text CONSTRAINT foo_bar_unique UNIQUE;
+    "#;
+        assert_debug_snapshot!(lint_sql(sql));
+    }
+    #[test]
+    fn test_unique_constraint_inline_add_column_unique() {
+        let sql = r#"
+ALTER TABLE foo ADD COLUMN bar text UNIQUE;
+"#;
+        assert_debug_snapshot!(lint_sql(sql));
+    }
 }
diff --git a/linter/src/rules/snapshots/squawk_linter__rules__disallow_unique_constraint__test_rules__unique_constraint_inline_add_column.snap b/linter/src/rules/snapshots/squawk_linter__rules__disallow_unique_constraint__test_rules__unique_constraint_inline_add_column.snap
new file mode 100644
index 00000000..29a903ed
--- /dev/null
+++ b/linter/src/rules/snapshots/squawk_linter__rules__disallow_unique_constraint__test_rules__unique_constraint_inline_add_column.snap
@@ -0,0 +1,23 @@
+---
+source: linter/src/rules/disallow_unique_constraint.rs
+expression: lint_sql(sql)
+---
+[
+    RuleViolation {
+        kind: DisallowedUniqueConstraint,
+        span: Span {
+            start: 0,
+            len: Some(
+                101,
+            ),
+        },
+        messages: [
+            Note(
+                "Adding a UNIQUE constraint requires an ACCESS EXCLUSIVE lock which blocks reads.",
+            ),
+            Help(
+                "Create an index CONCURRENTLY and create the constraint using the index.",
+            ),
+        ],
+    },
+]
diff --git a/linter/src/rules/snapshots/squawk_linter__rules__disallow_unique_constraint__test_rules__unique_constraint_inline_add_column_unique.snap b/linter/src/rules/snapshots/squawk_linter__rules__disallow_unique_constraint__test_rules__unique_constraint_inline_add_column_unique.snap
new file mode 100644
index 00000000..923caef5
--- /dev/null
+++ b/linter/src/rules/snapshots/squawk_linter__rules__disallow_unique_constraint__test_rules__unique_constraint_inline_add_column_unique.snap
@@ -0,0 +1,23 @@
+---
+source: linter/src/rules/disallow_unique_constraint.rs
+expression: lint_sql(sql)
+---
+[
+    RuleViolation {
+        kind: DisallowedUniqueConstraint,
+        span: Span {
+            start: 0,
+            len: Some(
+                47,
+            ),
+        },
+        messages: [
+            Note(
+                "Adding a UNIQUE constraint requires an ACCESS EXCLUSIVE lock which blocks reads.",
+            ),
+            Help(
+                "Create an index CONCURRENTLY and create the constraint using the index.",
+            ),
+        ],
+    },
+]