|
1 | 1 | """Alembic environment configuration.""" |
2 | 2 |
|
| 3 | +import asyncio |
3 | 4 | import os |
4 | 5 | from logging.config import fileConfig |
5 | 6 |
|
6 | | -from sqlalchemy import engine_from_config |
7 | | -from sqlalchemy import pool |
| 7 | +# Allow nested event loops (needed for pytest-asyncio and other async contexts) |
| 8 | +# Note: nest_asyncio doesn't work with uvloop, so we handle that case separately |
| 9 | +try: |
| 10 | + import nest_asyncio |
| 11 | + |
| 12 | + nest_asyncio.apply() |
| 13 | +except (ImportError, ValueError): |
| 14 | + # nest_asyncio not available or can't patch this loop type (e.g., uvloop) |
| 15 | + pass |
| 16 | + |
| 17 | +from sqlalchemy import engine_from_config, pool |
| 18 | +from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine |
8 | 19 |
|
9 | 20 | from alembic import context |
10 | 21 |
|
11 | | -from basic_memory.config import ConfigManager, DatabaseBackend |
| 22 | +from basic_memory.config import ConfigManager |
12 | 23 |
|
13 | 24 | # set config.env to "test" for pytest to prevent logging to file in utils.setup_logging() |
14 | 25 | os.environ["BASIC_MEMORY_ENV"] = "test" |
|
35 | 46 | sqlalchemy_url = DatabaseType.get_db_url( |
36 | 47 | app_config.database_path, DatabaseType.FILESYSTEM, app_config |
37 | 48 | ) |
38 | | - |
39 | | - # For Postgres, Alembic needs synchronous driver (psycopg2), not async (asyncpg) |
40 | | - if app_config.database_backend == DatabaseBackend.POSTGRES: |
41 | | - # Convert asyncpg URL to psycopg2 URL for Alembic |
42 | | - sqlalchemy_url = sqlalchemy_url.replace("postgresql+asyncpg://", "postgresql://") |
43 | | - |
44 | 49 | config.set_main_option("sqlalchemy.url", sqlalchemy_url) |
45 | 50 |
|
46 | 51 | # Interpret the config file for Python logging. |
@@ -85,28 +90,89 @@ def run_migrations_offline() -> None: |
85 | 90 | context.run_migrations() |
86 | 91 |
|
87 | 92 |
|
| 93 | +def do_run_migrations(connection): |
| 94 | + """Execute migrations with the given connection.""" |
| 95 | + context.configure( |
| 96 | + connection=connection, |
| 97 | + target_metadata=target_metadata, |
| 98 | + include_object=include_object, |
| 99 | + render_as_batch=True, |
| 100 | + compare_type=True, |
| 101 | + ) |
| 102 | + with context.begin_transaction(): |
| 103 | + context.run_migrations() |
| 104 | + |
| 105 | + |
| 106 | +async def run_async_migrations(connectable): |
| 107 | + """Run migrations asynchronously with AsyncEngine.""" |
| 108 | + async with connectable.connect() as connection: |
| 109 | + await connection.run_sync(do_run_migrations) |
| 110 | + await connectable.dispose() |
| 111 | + |
| 112 | + |
88 | 113 | def run_migrations_online() -> None: |
89 | 114 | """Run migrations in 'online' mode. |
90 | 115 |
|
91 | | - In this scenario we need to create an Engine |
92 | | - and associate a connection with the context. |
| 116 | + Supports both sync engines (SQLite) and async engines (PostgreSQL with asyncpg). |
93 | 117 | """ |
94 | | - connectable = engine_from_config( |
95 | | - config.get_section(config.config_ini_section, {}), |
96 | | - prefix="sqlalchemy.", |
97 | | - poolclass=pool.NullPool, |
98 | | - ) |
99 | | - |
100 | | - with connectable.connect() as connection: |
101 | | - context.configure( |
102 | | - connection=connection, |
103 | | - target_metadata=target_metadata, |
104 | | - include_object=include_object, |
105 | | - render_as_batch=True, |
106 | | - ) |
107 | | - |
108 | | - with context.begin_transaction(): |
109 | | - context.run_migrations() |
| 118 | + # Check if a connection/engine was provided (e.g., from run_migrations) |
| 119 | + connectable = context.config.attributes.get("connection", None) |
| 120 | + |
| 121 | + if connectable is None: |
| 122 | + # No connection provided, create engine from config |
| 123 | + url = context.config.get_main_option("sqlalchemy.url") |
| 124 | + |
| 125 | + # Check if it's an async URL (sqlite+aiosqlite or postgresql+asyncpg) |
| 126 | + if url and ("+asyncpg" in url or "+aiosqlite" in url): |
| 127 | + # Create async engine for asyncpg or aiosqlite |
| 128 | + connectable = create_async_engine( |
| 129 | + url, |
| 130 | + poolclass=pool.NullPool, |
| 131 | + future=True, |
| 132 | + ) |
| 133 | + else: |
| 134 | + # Create sync engine for regular sqlite or postgresql |
| 135 | + connectable = engine_from_config( |
| 136 | + context.config.get_section(context.config.config_ini_section, {}), |
| 137 | + prefix="sqlalchemy.", |
| 138 | + poolclass=pool.NullPool, |
| 139 | + ) |
| 140 | + |
| 141 | + # Handle async engines (PostgreSQL with asyncpg) |
| 142 | + if isinstance(connectable, AsyncEngine): |
| 143 | + # Try to run async migrations |
| 144 | + # nest_asyncio allows asyncio.run() from within event loops, but doesn't work with uvloop |
| 145 | + try: |
| 146 | + asyncio.run(run_async_migrations(connectable)) |
| 147 | + except RuntimeError as e: |
| 148 | + if "cannot be called from a running event loop" in str(e): |
| 149 | + # We're in a running event loop (likely uvloop) - need to use a different approach |
| 150 | + # Create a new thread to run the async migrations |
| 151 | + import concurrent.futures |
| 152 | + |
| 153 | + def run_in_thread(): |
| 154 | + """Run async migrations in a new event loop in a separate thread.""" |
| 155 | + new_loop = asyncio.new_event_loop() |
| 156 | + asyncio.set_event_loop(new_loop) |
| 157 | + try: |
| 158 | + new_loop.run_until_complete(run_async_migrations(connectable)) |
| 159 | + finally: |
| 160 | + new_loop.close() |
| 161 | + |
| 162 | + with concurrent.futures.ThreadPoolExecutor() as executor: |
| 163 | + future = executor.submit(run_in_thread) |
| 164 | + future.result() # Wait for completion and re-raise any exceptions |
| 165 | + else: |
| 166 | + raise |
| 167 | + else: |
| 168 | + # Handle sync engines (SQLite) or sync connections |
| 169 | + if hasattr(connectable, "connect"): |
| 170 | + # It's an engine, get a connection |
| 171 | + with connectable.connect() as connection: |
| 172 | + do_run_migrations(connection) |
| 173 | + else: |
| 174 | + # It's already a connection |
| 175 | + do_run_migrations(connectable) |
110 | 176 |
|
111 | 177 |
|
112 | 178 | if context.is_offline_mode(): |
|
0 commit comments