From 20fddc08f45ba33c08e301e7216ccbb227236609 Mon Sep 17 00:00:00 2001
From: ajspig <dragon@monstercode.com>
Date: Mon, 10 Aug 2026 12:03:46 -0400
Subject: [PATCH] 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>
---
 migrations/env.py | 66 ++++++++++++++++++++++++++---------------------
 1 file changed, 37 insertions(+), 29 deletions(-)

diff --git a/migrations/env.py b/migrations/env.py
index fabd6de7..2aed72af 100644
--- a/migrations/env.py
+++ b/migrations/env.py
@@ -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():