fix: dispose alembic engine after migrations run

engine_from_config in env.py leaked a pooled connection per alembic
invocation; repeated upgrades in the migration test suite accumulated
enough lingering sessions to block DROP DATABASE at fixture teardown.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
ajspig 2026-08-10 12:03:46 -04:00
parent 7a8880bafe
commit 20fddc08f4
1 changed files with 37 additions and 29 deletions

View File

@ -6,6 +6,7 @@ from urllib.parse import urlparse, urlunparse
from alembic import context
from sqlalchemy import engine_from_config, text
from sqlalchemy.engine import Connection
from src.config import settings
@ -150,37 +151,44 @@ def run_migrations_online() -> None:
},
)
with connectable.connect() as connection:
# Create schema and commit it outside the main migration transaction
connection.execute(
text(f"CREATE SCHEMA IF NOT EXISTS {target_metadata.schema};")
)
connection.execute(
text(f"GRANT ALL ON SCHEMA {target_metadata.schema} TO current_user")
)
# Install pgvector extension if it doesn't exist
connection.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
# Set and verify search_path
connection.execute(
text(f"SET search_path TO {target_metadata.schema}, public, extensions")
)
connection.commit()
try:
with connectable.connect() as connection:
_run_migrations_with_connection(connection)
finally:
# Release pooled connections immediately; lingering ones block
# DROP DATABASE in test harnesses that invoke alembic repeatedly.
connectable.dispose()
context.configure(
connection=connection,
target_metadata=target_metadata,
version_table_schema=target_metadata.schema,
include_schemas=True,
include_object=lambda obj, name, type_, reflected, compare_to: (
# Only include objects from our target schema
getattr(obj, "schema", None) == target_metadata.schema
if hasattr(obj, "schema")
else True
),
)
with context.begin_transaction():
context.run_migrations()
def _run_migrations_with_connection(connection: Connection) -> None:
# Create schema and commit it outside the main migration transaction
connection.execute(text(f"CREATE SCHEMA IF NOT EXISTS {target_metadata.schema};"))
connection.execute(
text(f"GRANT ALL ON SCHEMA {target_metadata.schema} TO current_user")
)
# Install pgvector extension if it doesn't exist
connection.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
# Set and verify search_path
connection.execute(
text(f"SET search_path TO {target_metadata.schema}, public, extensions")
)
connection.commit()
context.configure(
connection=connection,
target_metadata=target_metadata,
version_table_schema=target_metadata.schema,
include_schemas=True,
include_object=lambda obj, name, type_, reflected, compare_to: (
# Only include objects from our target schema
getattr(obj, "schema", None) == target_metadata.schema
if hasattr(obj, "schema")
else True
),
)
with context.begin_transaction():
context.run_migrations()
if context.is_offline_mode():