1212from alembic import context
1313from sqlalchemy import engine_from_config , pool
1414
15+ # Default Postgres timeouts applied to every migration. They keep a stuck
16+ # migration from queueing behind active writes and holding locks indefinitely.
17+ #
18+ # - lock_timeout: how long a statement waits for a lock before aborting. 3s
19+ # means a migration that cannot acquire its lock quickly gives up instead of
20+ # blocking writers behind it.
21+ # - statement_timeout: maximum runtime for any single statement. 30s catches
22+ # runaway DDL/UPDATEs; long index builds must use CREATE INDEX CONCURRENTLY
23+ # in an autocommit_block, which runs outside the transaction-bound timeout.
24+ # - idle_in_transaction_session_timeout: kills a transaction that has gone
25+ # idle while still holding locks (e.g. a stalled AccessExclusiveLock).
26+ #
27+ # These are session-level so they persist across each per-migration
28+ # transaction and across autocommit_block boundaries on the same connection.
29+ # Migration authors must NOT override them with `SET lock_timeout` or
30+ # `SET statement_timeout` inside a migration file — the migration linter
31+ # (scripts/ci_tools/migration_lint.py) flags those, with the
32+ # `migration-unsafe-ack` PR label as the documented escape hatch for
33+ # genuinely-long migrations that need a maintenance window.
34+ DEFAULT_MIGRATION_TIMEOUTS : dict [str , str ] = {
35+ "lock_timeout" : "3s" ,
36+ "statement_timeout" : "30s" ,
37+ "idle_in_transaction_session_timeout" : "10s" ,
38+ }
39+
40+
41+ def _format_set_statements (timeouts : dict [str , str ]) -> list [str ]:
42+ return [f"SET { key } { value } for key , value in timeouts .items ()]
43+
1544# Add explicit error handling to catch import errors
1645try :
1746 print ("Starting migration - importing modules" )
@@ -83,6 +112,8 @@ def run_migrations_offline() -> None:
83112 )
84113
85114 with context .begin_transaction ():
115+ for stmt in _format_set_statements (DEFAULT_MIGRATION_TIMEOUTS ):
116+ context .execute (stmt )
86117 context .run_migrations ()
87118 except Exception as e :
88119 print ("ERROR IN OFFLINE MIGRATIONS:" , str (e ))
@@ -106,7 +137,30 @@ def run_migrations_online() -> None:
106137 )
107138
108139 with connectable .connect () as connection :
109- context .configure (connection = connection , target_metadata = target_metadata )
140+ # Apply default migration timeouts at the session level so they
141+ # persist across per-migration transactions and any autocommit_block
142+ # boundaries opened by migrations (e.g. for CREATE INDEX CONCURRENTLY).
143+ #
144+ # exec_driver_sql autobegins a SQLAlchemy transaction. We commit it
145+ # before configure() so alembic doesn't latch onto it as an
146+ # "external" transaction — that mode disables transaction_per_migration
147+ # and breaks autocommit_block (which asserts self._transaction is not
148+ # None). Postgres SET is session-level, so the timeouts persist past
149+ # the commit.
150+ for stmt in _format_set_statements (DEFAULT_MIGRATION_TIMEOUTS ):
151+ connection .exec_driver_sql (stmt )
152+ connection .commit ()
153+
154+ # transaction_per_migration=True wraps each migration in its own
155+ # transaction (instead of a single outer transaction for all
156+ # migrations). This lets individual migrations opt into
157+ # autocommit_block() for operations that cannot run inside a
158+ # transaction, such as CREATE INDEX CONCURRENTLY.
159+ context .configure (
160+ connection = connection ,
161+ target_metadata = target_metadata ,
162+ transaction_per_migration = True ,
163+ )
110164
111165 with context .begin_transaction ():
112166 context .run_migrations ()
0 commit comments