From eea2049a4f2f34de88c6de781c1cafaab126c2a0 Mon Sep 17 00:00:00 2001 From: Dylan Griffith Date: Fri, 28 Aug 2026 14:03:27 +0200 Subject: [PATCH 1/3] Make prefer-robust-stmts aware of alembic version update transactions Today with Alembic you can write a migration like: ```py def upgrade() -> None: op.add_column('t', sa.Column('test_column', sa.Text(), nullable=True)) with op.get_context().autocommit_block(): op.execute("SELECT 1") def downgrade() -> None: """Downgrade schema.""" pass ``` This results in the following SQL: ```sql BEGIN; ALTER TABLE t ADD COLUMN test_column TEXT; COMMIT; SELECT 1; BEGIN; UPDATE alembic_version SET version_num='bbc810e82b46' WHERE alembic_version.version_num = 'b4d17e0c93aa'; COMMIT; ``` And it passes the "prefer-robust-stmts" check. But it shouldn't because as soon as you add this `autocommit_block` anywhere in your migration you are effectively breaking the wrapping migration for the whole alembic migration. And slightly surprisingly Alembic still decides to put a transaction at the start and the end of the migration. This is quite weird behaviour from Alembic. But it is resulting in this prefer-robust-stmts lint rule assuming that everything is "OK" because it sees that the `ALTER TABLE t` is in "some transaction". But it happens to be the wrong transaction and can definitely be partially failed. This PR makes a targetted implementation for alembic that looks more closely to work out whether or not the statement is actually in the same statement as the `alembic_version` update which is our main safety against partial failures. For now this code only handles alembic but if similar quirky behaviour is possible with other migration frameworks (or we want to protect against users putting explicitl `BEGIN ... COMMIT` in their migrations) then we might want to extend this implementation to find the relevant markers from the other frameworks. This PR also removes a comment about `IF NOT EXISTS` not being supported by alembic but this is definitely out of date as you can see from https://alembic.sqlalchemy.org/en/latest/ops.html#alembic.operations.Operations.add_column . --- .../src/rules/prefer_robust_stmts.rs | 161 ++++++++++++++++-- ...dd_column_in_separate_transaction_err.snap | 11 ++ ...ate_index_in_separate_transaction_err.snap | 13 ++ ...rop_table_in_separate_transaction_err.snap | 11 ++ docs/docs/prefer-robust-stmts.md | 4 - 5 files changed, 182 insertions(+), 18 deletions(-) create mode 100644 crates/squawk_linter/src/rules/snapshots/squawk_linter__rules__prefer_robust_stmts__test__alembic_add_column_in_separate_transaction_err.snap create mode 100644 crates/squawk_linter/src/rules/snapshots/squawk_linter__rules__prefer_robust_stmts__test__alembic_create_index_in_separate_transaction_err.snap create mode 100644 crates/squawk_linter/src/rules/snapshots/squawk_linter__rules__prefer_robust_stmts__test__alembic_drop_table_in_separate_transaction_err.snap diff --git a/crates/squawk_linter/src/rules/prefer_robust_stmts.rs b/crates/squawk_linter/src/rules/prefer_robust_stmts.rs index 9a4ced87c..f679a5d64 100644 --- a/crates/squawk_linter/src/rules/prefer_robust_stmts.rs +++ b/crates/squawk_linter/src/rules/prefer_robust_stmts.rs @@ -1,4 +1,4 @@ -use rustc_hash::FxHashMap; +use rustc_hash::{FxHashMap, FxHashSet}; use squawk_syntax::{ Parse, SourceFile, @@ -13,18 +13,58 @@ enum Constraint { Added, } +fn is_alembic_version_update(stmt: &ast::Stmt) -> bool { + matches!(stmt, ast::Stmt::Update(u) if u + .relation_name() + .and_then(|rn| rn.relation_name_ref()) + .and_then(|rnr| rnr.path_ref()) + .and_then(|path| path.segment()) + .is_some_and(|name| name.text().eq_ignore_ascii_case("alembic_version"))) +} + +fn alembic_version_commit_tx_indices(stmts: &[ast::Stmt]) -> Option> { + if !stmts.iter().any(is_alembic_version_update) { + return None; + } + + let mut indices = FxHashSet::default(); + let mut tx_start = None; + for (i, stmt) in stmts.iter().enumerate() { + match stmt { + ast::Stmt::Begin(_) => tx_start = Some(i), + ast::Stmt::Commit(_) | ast::Stmt::Rollback(_) => { + if let Some(start) = tx_start.take() + && stmts[start..=i].iter().any(is_alembic_version_update) + { + indices.extend(start..=i); + } + } + _ => {} + } + } + Some(indices) +} + pub(crate) fn prefer_robust_stmts(ctx: &mut Linter, parse: &Parse) { let file = parse.tree(); + let stmts: Vec = file.stmts().collect(); let mut inside_transaction = ctx.settings.assume_in_transaction; let mut constraint_names: FxHashMap = FxHashMap::default(); + let version_tx_indices = alembic_version_commit_tx_indices(&stmts); + enum ActionErrorMessage { IfExists, IfNotExists, None, } - for stmt in file.stmts() { + for (i, stmt) in stmts.iter().enumerate() { + let in_robust_tx = match &version_tx_indices { + Some(set) => ctx.settings.assume_in_transaction || set.contains(&i), + None => inside_transaction, + }; + match stmt { ast::Stmt::Begin(_) => { inside_transaction = true; @@ -107,7 +147,7 @@ pub(crate) fn prefer_robust_stmts(ctx: &mut Linter, parse: &Parse) { _ => (ActionErrorMessage::None, None), }; - if inside_transaction { + if in_robust_tx { continue; } @@ -132,7 +172,7 @@ pub(crate) fn prefer_robust_stmts(ctx: &mut Linter, parse: &Parse) { ast::Stmt::CreateIndex(create_index) if create_index.if_not_exists().is_none() && create_index.index().is_some() - && (create_index.concurrently_token().is_some() || !inside_transaction) => + && (create_index.concurrently_token().is_some() || !in_robust_tx) => { let fix = create_index.index().map(|index| { let at = index.syntax().text_range().start(); @@ -146,7 +186,7 @@ pub(crate) fn prefer_robust_stmts(ctx: &mut Linter, parse: &Parse) { ).help("Use an explicit name for a concurrently created index").fix(fix)); } ast::Stmt::CreateTable(create_table) - if create_table.if_not_exists().is_none() && !inside_transaction => + if create_table.if_not_exists().is_none() && !in_robust_tx => { let is_temp = create_table .persistence() @@ -173,7 +213,7 @@ pub(crate) fn prefer_robust_stmts(ctx: &mut Linter, parse: &Parse) { ).fix(fix)); } ast::Stmt::DropIndex(drop_index) - if drop_index.if_exists().is_none() && !inside_transaction => + if drop_index.if_exists().is_none() && !in_robust_tx => { let fix = drop_index.index_refs().next().map(|first_index| { let at = first_index.syntax().text_range().start(); @@ -188,7 +228,7 @@ pub(crate) fn prefer_robust_stmts(ctx: &mut Linter, parse: &Parse) { ).fix(fix)); } ast::Stmt::DropTable(drop_table) - if drop_table.if_exists().is_none() && !inside_transaction => + if drop_table.if_exists().is_none() && !in_robust_tx => { let fix = drop_table.table_token().map(|table_token| { let at = table_token.text_range().end(); @@ -201,9 +241,7 @@ pub(crate) fn prefer_robust_stmts(ctx: &mut Linter, parse: &Parse) { drop_table.syntax(), ).fix(fix)); } - ast::Stmt::DropType(drop_type) - if drop_type.if_exists().is_none() && !inside_transaction => - { + ast::Stmt::DropType(drop_type) if drop_type.if_exists().is_none() && !in_robust_tx => { let fix = drop_type.type_token().map(|type_token| { let at = type_token.text_range().end(); let edit = Edit::insert(" if exists", at); @@ -245,6 +283,105 @@ mod test { crate::test_utils::lint_ok_with(sql, settings, Rule::PreferRobustStmts); } + #[test] + fn alembic_add_column_in_separate_transaction_err() { + let sql = r#" +BEGIN; +ALTER TABLE t ADD COLUMN test_column TEXT; +COMMIT; + +SELECT 1; + +BEGIN; +UPDATE alembic_version SET version_num='bbc810e82b46' WHERE alembic_version.version_num = 'b4d17e0c93aa'; +COMMIT; + "#; + assert_snapshot!(lint_errors(sql)); + } + + #[test] + fn alembic_add_column_with_if_not_exists_in_separate_transaction_ok() { + let sql = r#" +BEGIN; +ALTER TABLE t ADD COLUMN IF NOT EXISTS test_column TEXT; +COMMIT; + +BEGIN; +UPDATE alembic_version SET version_num='bbc810e82b46' WHERE alembic_version.version_num = 'b4d17e0c93aa'; +COMMIT; + "#; + lint_ok(sql); + } + + #[test] + fn alembic_add_column_in_version_transaction_ok() { + let sql = r#" +BEGIN; +ALTER TABLE t ADD COLUMN test_column TEXT; +UPDATE alembic_version SET version_num='bbc810e82b46' WHERE alembic_version.version_num = 'b4d17e0c93aa'; +COMMIT; + "#; + lint_ok(sql); + } + + #[test] + fn alembic_multiple_migrations_in_own_transactions_ok() { + let sql = r#" +BEGIN; +ALTER TABLE t ADD COLUMN a TEXT; +UPDATE alembic_version SET version_num='v2' WHERE version_num='v1'; +COMMIT; + +BEGIN; +ALTER TABLE t ADD COLUMN b TEXT; +UPDATE alembic_version SET version_num='v3' WHERE version_num='v2'; +COMMIT; + "#; + lint_ok(sql); + } + + #[test] + fn alembic_drop_table_in_separate_transaction_err() { + let sql = r#" +BEGIN; +DROP TABLE old_table; +COMMIT; + +BEGIN; +UPDATE alembic_version SET version_num='abc'; +COMMIT; + "#; + assert_snapshot!(lint_errors(sql)); + } + + #[test] + fn alembic_drop_table_with_if_exists_in_separate_transaction_ok() { + let sql = r#" +BEGIN; +DROP TABLE IF EXISTS old_table; +COMMIT; + +BEGIN; +UPDATE alembic_version SET version_num='abc'; +COMMIT; + "#; + lint_ok(sql); + } + + #[test] + fn alembic_create_index_in_separate_transaction_err() { + let sql = r#" +BEGIN; +CREATE INDEX foo_idx ON bar (baz); +COMMIT; + +BEGIN; +UPDATE alembic_version SET version_num='abc'; +COMMIT; + "#; + assert_snapshot!(lint_errors(sql)); + } + #[test] fn fix_drop_type_if_exists() { assert_snapshot!(fix(" @@ -415,7 +552,6 @@ CREATE TABLE IF NOT EXISTS "core_bar" ( #[test] fn prefer_robust_stmt_part_6_ok() { - // If done in a transaction, most forms of drop are fine let sql = r#" BEGIN; DROP INDEX "core_bar_foo_id_idx"; @@ -428,7 +564,6 @@ COMMIT; #[test] fn select_ok() { - // select is fine, we're only interested in modifications to the tables let sql = r#" select 1; -- so we don't skip checking SELECT 1; @@ -438,7 +573,6 @@ SELECT 1; #[test] fn insert_ok() { - // select is fine, we're only interested in modifications to the tables let sql = r#" select 1; -- so we don't skip checking INSERT INTO tbl VALUES (a); @@ -448,7 +582,6 @@ INSERT INTO tbl VALUES (a); #[test] fn alter_table_ok() { - // select is fine, we're only interested in modifications to the tables let sql = r#" select 1; -- so we don't skip checking ALTER TABLE "core_foo" DROP CONSTRAINT IF EXISTS "core_foo_idx"; diff --git a/crates/squawk_linter/src/rules/snapshots/squawk_linter__rules__prefer_robust_stmts__test__alembic_add_column_in_separate_transaction_err.snap b/crates/squawk_linter/src/rules/snapshots/squawk_linter__rules__prefer_robust_stmts__test__alembic_add_column_in_separate_transaction_err.snap new file mode 100644 index 000000000..3cb98294a --- /dev/null +++ b/crates/squawk_linter/src/rules/snapshots/squawk_linter__rules__prefer_robust_stmts__test__alembic_add_column_in_separate_transaction_err.snap @@ -0,0 +1,11 @@ +--- +source: crates/squawk_linter/src/rules/prefer_robust_stmts.rs +expression: lint_errors(sql) +--- +warning[prefer-robust-stmts]: Missing `IF NOT EXISTS`, the migration can't be rerun if it fails part way through. + ╭▸ +3 │ ALTER TABLE t ADD COLUMN test_column TEXT; + │ ━━━━━━━━━━━━━━━━━━━━━━━━━━━ + ╭╴ +3 │ ALTER TABLE t ADD COLUMN if not exists test_column TEXT; + ╰╴ +++++++++++++ diff --git a/crates/squawk_linter/src/rules/snapshots/squawk_linter__rules__prefer_robust_stmts__test__alembic_create_index_in_separate_transaction_err.snap b/crates/squawk_linter/src/rules/snapshots/squawk_linter__rules__prefer_robust_stmts__test__alembic_create_index_in_separate_transaction_err.snap new file mode 100644 index 000000000..2ab80b028 --- /dev/null +++ b/crates/squawk_linter/src/rules/snapshots/squawk_linter__rules__prefer_robust_stmts__test__alembic_create_index_in_separate_transaction_err.snap @@ -0,0 +1,13 @@ +--- +source: crates/squawk_linter/src/rules/prefer_robust_stmts.rs +expression: lint_errors(sql) +--- +warning[prefer-robust-stmts]: Missing `IF NOT EXISTS`, the migration can't be rerun if it fails part way through. + ╭▸ +3 │ CREATE INDEX foo_idx ON bar (baz); + │ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + │ + ├ help: Use an explicit name for a concurrently created index + ╭╴ +3 │ CREATE INDEX if not exists foo_idx ON bar (baz); + ╰╴ +++++++++++++ diff --git a/crates/squawk_linter/src/rules/snapshots/squawk_linter__rules__prefer_robust_stmts__test__alembic_drop_table_in_separate_transaction_err.snap b/crates/squawk_linter/src/rules/snapshots/squawk_linter__rules__prefer_robust_stmts__test__alembic_drop_table_in_separate_transaction_err.snap new file mode 100644 index 000000000..02232c773 --- /dev/null +++ b/crates/squawk_linter/src/rules/snapshots/squawk_linter__rules__prefer_robust_stmts__test__alembic_drop_table_in_separate_transaction_err.snap @@ -0,0 +1,11 @@ +--- +source: crates/squawk_linter/src/rules/prefer_robust_stmts.rs +expression: lint_errors(sql) +--- +warning[prefer-robust-stmts]: Missing `IF EXISTS`, the migration can't be rerun if it fails part way through. + ╭▸ +3 │ DROP TABLE old_table; + │ ━━━━━━━━━━━━━━━━━━━━━ + ╭╴ +3 │ DROP TABLE if exists old_table; + ╰╴ +++++++++ diff --git a/docs/docs/prefer-robust-stmts.md b/docs/docs/prefer-robust-stmts.md index f9c476a8b..de1828179 100644 --- a/docs/docs/prefer-robust-stmts.md +++ b/docs/docs/prefer-robust-stmts.md @@ -140,7 +140,3 @@ DROP INDEX "foo_idx"; -- use: DROP INDEX IF EXISTS "foo_idx"; ``` - - -## solution for alembic and sqlalchemy -Alembic doesn't support `IF NOT EXISTS`. You must use raw SQL via [`op.execute()`](https://alembic.sqlalchemy.org/en/latest/ops.html#alembic.operations.Operations.execute). See this [GitHub Issue tracking `IF NOT EXISTS` support](https://github.com/sqlalchemy/alembic/issues/151) for more information. From 6f2e1c9e56c89af3a832427b1a8aa3fddcefcf53 Mon Sep 17 00:00:00 2001 From: Dylan Griffith Date: Sat, 29 Aug 2026 09:05:48 +0200 Subject: [PATCH 2/3] Dont delete comments --- crates/squawk_linter/src/rules/prefer_robust_stmts.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/crates/squawk_linter/src/rules/prefer_robust_stmts.rs b/crates/squawk_linter/src/rules/prefer_robust_stmts.rs index f679a5d64..d5d08ee23 100644 --- a/crates/squawk_linter/src/rules/prefer_robust_stmts.rs +++ b/crates/squawk_linter/src/rules/prefer_robust_stmts.rs @@ -552,6 +552,7 @@ CREATE TABLE IF NOT EXISTS "core_bar" ( #[test] fn prefer_robust_stmt_part_6_ok() { + // If done in a transaction, most forms of drop are fine let sql = r#" BEGIN; DROP INDEX "core_bar_foo_id_idx"; @@ -564,6 +565,7 @@ COMMIT; #[test] fn select_ok() { + // select is fine, we're only interested in modifications to the tables let sql = r#" select 1; -- so we don't skip checking SELECT 1; @@ -573,6 +575,7 @@ SELECT 1; #[test] fn insert_ok() { + // select is fine, we're only interested in modifications to the tables let sql = r#" select 1; -- so we don't skip checking INSERT INTO tbl VALUES (a); @@ -582,6 +585,7 @@ INSERT INTO tbl VALUES (a); #[test] fn alter_table_ok() { + // select is fine, we're only interested in modifications to the tables let sql = r#" select 1; -- so we don't skip checking ALTER TABLE "core_foo" DROP CONSTRAINT IF EXISTS "core_foo_idx"; From 6ade42a48615db4968a2763ac9995189981f1acd Mon Sep 17 00:00:00 2001 From: Dylan Griffith Date: Sat, 29 Aug 2026 09:26:28 +0200 Subject: [PATCH 3/3] Only check is_alembic_version_update once --- crates/squawk_linter/src/rules/prefer_robust_stmts.rs | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/crates/squawk_linter/src/rules/prefer_robust_stmts.rs b/crates/squawk_linter/src/rules/prefer_robust_stmts.rs index d5d08ee23..8efe4ff4b 100644 --- a/crates/squawk_linter/src/rules/prefer_robust_stmts.rs +++ b/crates/squawk_linter/src/rules/prefer_robust_stmts.rs @@ -23,10 +23,6 @@ fn is_alembic_version_update(stmt: &ast::Stmt) -> bool { } fn alembic_version_commit_tx_indices(stmts: &[ast::Stmt]) -> Option> { - if !stmts.iter().any(is_alembic_version_update) { - return None; - } - let mut indices = FxHashSet::default(); let mut tx_start = None; for (i, stmt) in stmts.iter().enumerate() { @@ -42,7 +38,7 @@ fn alembic_version_commit_tx_indices(stmts: &[ast::Stmt]) -> Option {} } } - Some(indices) + (!indices.is_empty()).then_some(indices) } pub(crate) fn prefer_robust_stmts(ctx: &mut Linter, parse: &Parse) {