"""Alembic environment for DeerFlow application tables. ONLY manages DeerFlow's tables (runs, threads_meta, cron_jobs, users). LangGraph's checkpointer tables are managed by LangGraph itself -- they have their own schema lifecycle and must not be touched by Alembic. """ from __future__ import annotations import asyncio import logging import os from logging.config import fileConfig from alembic import context from sqlalchemy import text from sqlalchemy.ext.asyncio import create_async_engine from deerflow.persistence.base import Base # Import all models so metadata is populated. try: import deerflow.persistence.models as models # register ORM models with Base.metadata _ = models except ImportError: # Models not available — migration will work with existing metadata only. logging.getLogger(__name__).warning("Could not import deerflow.persistence.models; Alembic may not detect all tables") config = context.config if config.config_file_name is not None: fileConfig(config.config_file_name) target_metadata = Base.metadata def run_migrations_offline() -> None: url = config.get_main_option("sqlalchemy.url") context.configure( url=url, target_metadata=target_metadata, literal_binds=True, render_as_batch=True, ) with context.begin_transaction(): context.run_migrations() def do_run_migrations(connection): render_as_batch = connection.dialect.name == "sqlite" # On MySQL, shrink the session lock waits so a blocked DDL fails fast # instead of hanging on the server default (lock_wait_timeout = 31536000s # = 1 year; innodb_lock_wait_timeout = 50s). Without this, an idle # transaction on another connection can freeze startup indefinitely. # Tunable via DEER_FLOW_ALEMBIC_LOCK_WAIT (seconds, default 15). if connection.dialect.name == "mysql": try: lock_wait = int(os.getenv("DEER_FLOW_ALEMBIC_LOCK_WAIT", "15")) except ValueError: lock_wait = 15 connection.execute(text(f"SET SESSION lock_wait_timeout = {lock_wait}")) connection.execute(text(f"SET SESSION innodb_lock_wait_timeout = {lock_wait}")) context.configure( connection=connection, target_metadata=target_metadata, render_as_batch=render_as_batch, # Required for SQLite ALTER TABLE support ) with context.begin_transaction(): context.run_migrations() async def run_migrations_online() -> None: connectable = create_async_engine(config.get_main_option("sqlalchemy.url")) async with connectable.connect() as connection: await connection.run_sync(do_run_migrations) await connectable.dispose() if context.is_offline_mode(): run_migrations_offline() else: asyncio.run(run_migrations_online())