linter: add file-level override for assume_in_transaction

crates/squawk_linter/src/ignore.rs

@@ -19,7 +19,7 @@ pub struct Ignore {
     pub kind: IgnoreKind,
 }
 
-fn comment_body(token: &SyntaxToken) -> Option<(&str, TextRange)> {
+pub(crate) fn comment_body(token: &SyntaxToken) -> Option<(&str, TextRange)> {
     let range = token.text_range();
     if token.kind() == SyntaxKind::COMMENT {
         let text = token.text();
@@ -128,6 +128,34 @@ pub(crate) fn find_ignores(ctx: &mut Linter, file: &SyntaxNode) {
     }
 }
 
+const DISABLE_ASSUME_IN_TRANSACTION: &str = "squawk-disable-assume-in-transaction";
+const ENABLE_ASSUME_IN_TRANSACTION: &str = "squawk-enable-assume-in-transaction";
+
+pub fn find_transaction_override(file: &SyntaxNode) -> Option<bool> {
+    for event in file.preorder_with_tokens() {
+        match event {
+            rowan::WalkEvent::Enter(NodeOrToken::Token(token))
+                if token.kind() == SyntaxKind::COMMENT =>
+            {
+                if let Some((body, _range)) = comment_body(&token) {
+                    let trimmed = body.trim();
+                    let trimmed = trimmed
+                        .find("--")
+                        .map_or(trimmed, |idx| trimmed[..idx].trim_end());
+                    if trimmed == DISABLE_ASSUME_IN_TRANSACTION {
+                        return Some(false);
+                    }
+                    if trimmed == ENABLE_ASSUME_IN_TRANSACTION {
+                        return Some(true);
+                    }
+                }
+            }
+            _ => (),
+        }
+    }
+    None
+}
+
 #[cfg(test)]
 mod test {
 
@@ -612,4 +640,44 @@ alter table t drop column c cascade;
         ]
         ");
     }
+
+    #[test]
+    fn disable_assume_in_transaction() {
+        use super::find_transaction_override;
+        let sql = "-- squawk-disable-assume-in-transaction\nSELECT 1;";
+        let parse = squawk_syntax::SourceFile::parse(sql);
+        assert_eq!(find_transaction_override(&parse.syntax_node()), Some(false));
+    }
+
+    #[test]
+    fn enable_assume_in_transaction() {
+        use super::find_transaction_override;
+        let sql = "-- squawk-enable-assume-in-transaction\nSELECT 1;";
+        let parse = squawk_syntax::SourceFile::parse(sql);
+        assert_eq!(find_transaction_override(&parse.syntax_node()), Some(true));
+    }
+
+    #[test]
+    fn disable_assume_in_transaction_c_style_comment() {
+        use super::find_transaction_override;
+        let sql = "/* squawk-disable-assume-in-transaction */\nSELECT 1;";
+        let parse = squawk_syntax::SourceFile::parse(sql);
+        assert_eq!(find_transaction_override(&parse.syntax_node()), Some(false));
+    }
+
+    #[test]
+    fn disable_assume_in_transaction_with_trailing_comment() {
+        use super::find_transaction_override;
+        let sql = "-- squawk-disable-assume-in-transaction -- not in a transaction\nSELECT 1;";
+        let parse = squawk_syntax::SourceFile::parse(sql);
+        assert_eq!(find_transaction_override(&parse.syntax_node()), Some(false));
+    }
+
+    #[test]
+    fn transaction_override_none_when_absent() {
+        use super::find_transaction_override;
+        let sql = "SELECT 1;";
+        let parse = squawk_syntax::SourceFile::parse(sql);
+        assert_eq!(find_transaction_override(&parse.syntax_node()), None);
+    }
 }

crates/squawk_linter/src/lib.rs

@@ -5,6 +5,7 @@ use enum_iterator::Sequence;
 use enum_iterator::all;
 pub use ignore::Ignore;
 use ignore::find_ignores;
+use ignore::find_transaction_override;
 use ignore_index::IgnoreIndex;
 use rowan::TextRange;
 use rowan::TextSize;
@@ -332,6 +333,10 @@ impl Linter {
 
     #[must_use]
     pub fn lint(&mut self, file: &Parse<SourceFile>, text: &str) -> Vec<Violation> {
+        if let Some(override_value) = find_transaction_override(&file.syntax_node()) {
+            self.settings.assume_in_transaction = override_value;
+        }
+
         if self.rules.contains(&Rule::AddingFieldWithDefault) {
             adding_field_with_default(self, file);
         }

crates/squawk_linter/src/rules/ban_concurrent_index_creation_in_transaction.rs

@@ -134,4 +134,20 @@ mod test {
             },
         );
     }
+
+    #[test]
+    fn squawk_disable_assume_in_transaction_overrides() {
+        let sql = r#"
+-- squawk-disable-assume-in-transaction
+CREATE INDEX CONCURRENTLY "field_name_idx" ON "table_name" ("field_name");
+ALTER TABLE "table_name" ADD CONSTRAINT "field_name_id" UNIQUE USING INDEX "field_name_idx";
+    "#;
+        lint_ok_with(
+            sql,
+            LinterSettings {
+                assume_in_transaction: true,
+                ..Default::default()
+            },
+        );
+    }
 }

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

@@ -0,0 +1,17 @@
+---
+source: crates/squawk_linter/src/rules/transaction_nesting.rs
+assertion_line: 195
+expression: "lint_errors_with(sql, LinterSettings\n{ assume_in_transaction: false, ..Default::default() },)"
+---
+warning[transaction-nesting]: There is an existing transaction already in progress, managed by your migration tool.
+  ╭▸ 
+3 │ BEGIN;
+  │ ━━━━━
+  │
+  ╰ help: Put migration statements in separate files to have them be in separate transactions or don't use the assume-in-transaction setting.
+warning[transaction-nesting]: Attempting to end the transaction that is managed by your migration tool
+  ╭▸ 
+5 │ COMMIT;
+  │ ━━━━━━
+  │
+  ╰ help: Put migration statements in separate files to have them be in separate transactions or don't use the assume-in-transaction setting.

crates/squawk_linter/src/rules/transaction_nesting.rs

@@ -168,4 +168,36 @@ SELECT 1;
         };
         lint_ok_with(sql, settings);
     }
+
+    #[test]
+    fn squawk_disable_assume_in_transaction_allows_begin_commit() {
+        let sql = r#"
+-- squawk-disable-assume-in-transaction
+BEGIN;
+SELECT 1;
+COMMIT;
+        "#;
+        let settings = LinterSettings {
+            assume_in_transaction: true,
+            ..Default::default()
+        };
+        lint_ok_with(sql, settings);
+    }
+
+    #[test]
+    fn squawk_enable_assume_in_transaction_flags_begin() {
+        let sql = r#"
+-- squawk-enable-assume-in-transaction
+BEGIN;
+SELECT 1;
+COMMIT;
+        "#;
+        assert_snapshot!(lint_errors_with(
+            sql,
+            LinterSettings {
+                assume_in_transaction: false,
+                ..Default::default()
+            },
+        ));
+    }
 }

docs/docs/cli.md

@@ -54,6 +54,22 @@ alter table t drop column c cascade;
 create table t (a int);
 ```
 
+### Overriding `assume_in_transaction` via comments
+
+If you have `assume_in_transaction = true` set globally (via config or CLI flag), you can disable it for a specific file with `squawk-disable-assume-in-transaction`:
+
+```sql
+-- squawk-disable-assume-in-transaction
+CREATE INDEX CONCURRENTLY IF NOT EXISTS my_idx ON my_table (col);
+```
+
+Similarly, you can enable it for a specific file with `squawk-enable-assume-in-transaction`:
+
+```sql
+-- squawk-enable-assume-in-transaction
+ALTER TABLE my_table ADD COLUMN my_col integer;
+```
+
 ## Files
 
 Files can be excluded from linting via the `--exclude-path` flag. Glob matching is supported and the flag can be provided multiple times.