deerflow-code/offline-backend-20260512/backend/packages/harness/deerflow/persistence/engine.py

"""Async SQLAlchemy engine lifecycle management.

Initializes at Gateway startup, provides session factory for
repositories, disposes at shutdown.

When database.backend="memory", init_engine is a no-op and
get_session_factory() returns None. Repositories must check for
None and fall back to in-memory implementations.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import sys
from importlib.util import find_spec
from pathlib import Path

from sqlalchemy import text
from sqlalchemy.engine import make_url
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine


def _json_serializer(obj: object) -> str:
    """JSON serializer with ensure_ascii=False for Chinese character support."""
    return json.dumps(obj, ensure_ascii=False)


logger = logging.getLogger(__name__)

_engine: AsyncEngine | None = None
_session_factory: async_sessionmaker[AsyncSession] | None = None


def _add_runtime_vendor_path() -> None:
    """Allow offline deployments to place optional DB drivers in DEER_FLOW_HOME.

    The offline Docker image keeps the Python environment inside the image, while
    ``/app/backend/.deer-flow`` is a persistent bind mount.  The MySQL switch
    script installs ``asyncmy`` into ``{DEER_FLOW_HOME}/python-vendor`` so the
    unchanged startup script can still import it after container recreation.
    """
    try:
        from deerflow.config.runtime_paths import runtime_home

        vendor_dir = runtime_home() / "python-vendor"
    except Exception:
        vendor_dir = Path(".deer-flow/python-vendor").resolve()

    vendor = str(vendor_dir)
    if vendor_dir.exists() and vendor not in sys.path:
        sys.path.insert(0, vendor)


def _mysql_driver_name(url: str) -> str:
    if url.startswith("mysql+aiomysql://"):
        return "aiomysql"
    return "asyncmy"


async def _run_alembic_upgrade(url: str) -> None:
    """Apply pending Alembic migrations against ``url`` (``alembic upgrade head``).

    Runs after ``Base.metadata.create_all`` so brand-new deployments still get
    a complete schema even without Alembic, while existing deployments pick up
    in-place schema fixes (column widenings, new indexes, etc.) at the next
    restart — no manual ``alembic upgrade`` step required.

    Notes:

    * Each historical migration in ``migrations/versions/`` already guards
      itself with ``_has_table`` / ``_has_column``, so re-running them against
      a database that was originally built via ``create_all`` is a sequence of
      no-ops up to the first migration that actually has work to do.
    * The Alembic ``env.py`` calls ``asyncio.run`` internally, which cannot run
      inside an already-running loop — so we hop into a worker thread.
    * Failures here are logged but do not abort startup. The app remains
      operable on the ``create_all`` schema; the operator just won't get
      in-place fixes until they investigate.
    """
    if os.getenv("DEER_FLOW_SKIP_ALEMBIC", "").lower() in ("1", "true", "yes"):
        logger.warning("DEER_FLOW_SKIP_ALEMBIC is set; skipping Alembic migrations (schema stays at create_all)")
        return

    ini_path = Path(__file__).parent / "migrations" / "alembic.ini"
    if not ini_path.exists():
        logger.warning("Alembic config not found at %s; skipping auto-migrate", ini_path)
        return

    # Alembic's ``Config`` is backed by ConfigParser, which performs ``%``
    # interpolation on every option value. URL-encoded credentials such as
    # ``kland%4088124932`` (the ``@`` in ``kland@88124932``) trip that up with
    # ``ValueError: invalid interpolation syntax ... at position N``. Doubling
    # the percent signs neutralises the interpolation; ConfigParser un-doubles
    # them back to the original value when read.
    escaped_url = url.replace("%", "%%")

    def _do_upgrade() -> None:
        # Import inside the worker so a missing Alembic install can't break
        # the engine import chain.
        from alembic import command
        from alembic.config import Config

        cfg = Config(str(ini_path))
        cfg.set_main_option("sqlalchemy.url", escaped_url)
        # The ``script_location`` in alembic.ini is ``%(here)s`` which points
        # at the directory of alembic.ini itself — already correct for us.
        command.upgrade(cfg, "head")

    # Bound the upgrade so a stuck DDL (e.g. MySQL metadata-lock wait) cannot
    # hang startup forever. Tunable via DEER_FLOW_ALEMBIC_TIMEOUT (seconds).
    # NOTE: asyncio.wait_for cancels the *await* on timeout; the worker thread
    # itself keeps running (Python can't kill threads). That's acceptable —
    # the orphan thread will exit when its MySQL connection eventually errors
    # or when the process restarts. The point here is to unblock startup.
    try:
        timeout_s = float(os.getenv("DEER_FLOW_ALEMBIC_TIMEOUT", "60"))
    except ValueError:
        timeout_s = 60.0

    try:
        await asyncio.wait_for(asyncio.to_thread(_do_upgrade), timeout=timeout_s)
        logger.info("Alembic migrations applied (or already up-to-date)")
    except asyncio.TimeoutError:
        logger.error(
            "Alembic upgrade did not finish within %.0fs; continuing startup on create_all schema. "
            "Likely cause: MySQL metadata-lock wait — check SHOW FULL PROCESSLIST for a blocking transaction.",
            timeout_s,
        )
    except Exception:
        logger.exception("Auto-applying Alembic migrations failed; continuing with create_all schema")


async def _auto_create_postgres_db(url: str) -> None:
    """Connect to the ``postgres`` maintenance DB and CREATE DATABASE.

    The target database name is extracted from *url*.  The connection is
    made to the default ``postgres`` database on the same server using
    ``AUTOCOMMIT`` isolation (CREATE DATABASE cannot run inside a
    transaction).
    """
    from sqlalchemy import text
    from sqlalchemy.engine.url import make_url

    parsed = make_url(url)
    db_name = parsed.database
    if not db_name:
        raise ValueError("Cannot auto-create database: no database name in URL")

    # Connect to the default 'postgres' database to issue CREATE DATABASE
    maint_url = parsed.set(database="postgres")
    maint_engine = create_async_engine(maint_url, isolation_level="AUTOCOMMIT")
    try:
        async with maint_engine.connect() as conn:
            await conn.execute(text(f'CREATE DATABASE "{db_name}"'))
        logger.info("Auto-created PostgreSQL database: %s", db_name)
    finally:
        await maint_engine.dispose()


async def _auto_create_mysql_db(url: str) -> None:
    """Create the target MySQL database when it is missing.

    ``create_all`` can create tables, but MySQL requires the database/schema to
    exist before the application URL can connect. We connect to the built-in
    ``mysql`` maintenance database on the same server and issue a conservative
    ``CREATE DATABASE IF NOT EXISTS``.
    """

    parsed = make_url(url)
    db_name = parsed.database
    if not db_name:
        raise ValueError("Cannot auto-create MySQL database: no database name in URL")

    maint_url = parsed.set(database="mysql")
    maint_engine = create_async_engine(
        maint_url.render_as_string(hide_password=False),
        isolation_level="AUTOCOMMIT",
        pool_pre_ping=True,
    )
    quoted = "`" + db_name.replace("`", "``") + "`"
    try:
        async with maint_engine.connect() as conn:
            await conn.execute(
                text(
                    f"CREATE DATABASE IF NOT EXISTS {quoted} "
                    "CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"
                )
            )
        logger.info("Auto-created MySQL database if missing: %s", db_name)
    finally:
        await maint_engine.dispose()


def _ensure_orm_columns_sync(sync_conn) -> None:
    """Add ORM-declared columns that are missing from already-existing tables.

    ``Base.metadata.create_all`` creates brand-new tables but never patches
    existing ones, and Alembic auto-upgrade is disabled (see the commented-out
    ``_run_alembic_upgrade`` call) — so on an existing database a new
    ``mapped_column`` on an ORM model would otherwise error at first query
    (e.g. ``no such column: roundtable_drafts.task_id`` on SQLite).

    This sweep compares each existing table against the ORM metadata and issues
    ``ALTER TABLE … ADD COLUMN`` for every missing column that can be added
    safely (nullable, or carrying a server_default). Indexes covering a newly
    added column are created as well. Works on sqlite/mysql/postgres — the
    column DDL is compiled through the active dialect, so portable types
    (``PortableJSON``/``PortableLongText``/``BeijingDateTime``) resolve to the
    right backend type. Everything is per-column best-effort: a failure is
    logged and startup continues.
    """
    from sqlalchemy import inspect as sa_inspect
    from sqlalchemy.schema import CreateColumn

    from deerflow.persistence.base import Base

    try:
        import deerflow.persistence.models  # noqa: F401 — register all ORM models on Base.metadata
    except ImportError:
        pass

    inspector = sa_inspect(sync_conn)
    existing_tables = set(inspector.get_table_names())
    preparer = sync_conn.dialect.identifier_preparer

    for table in Base.metadata.sorted_tables:
        if table.name not in existing_tables:
            continue
        existing_cols = {c["name"] for c in inspector.get_columns(table.name)}
        missing = [c for c in table.columns if c.name not in existing_cols]
        if not missing:
            continue

        added: set[str] = set()
        for col in missing:
            if not col.nullable and col.server_default is None:
                logger.warning(
                    "Cannot auto-add NOT NULL column without server_default: %s.%s — add it manually",
                    table.name,
                    col.name,
                )
                continue
            try:
                spec = CreateColumn(col).compile(dialect=sync_conn.dialect).string
                sync_conn.exec_driver_sql(f"ALTER TABLE {preparer.quote(table.name)} ADD COLUMN {spec}")
                added.add(col.name)
                logger.info("Auto-added missing column %s.%s", table.name, col.name)
            except Exception:
                logger.exception("Failed to auto-add column %s.%s", table.name, col.name)

        if not added:
            continue
        try:
            existing_idx = {ix["name"] for ix in inspector.get_indexes(table.name)}
        except Exception:
            existing_idx = set()
        for idx in table.indexes:
            if idx.name in existing_idx or not any(c.name in added for c in idx.columns):
                continue
            try:
                idx.create(sync_conn)
                logger.info("Auto-created index %s on %s", idx.name, table.name)
            except Exception:
                logger.exception("Failed to auto-create index %s on %s", idx.name, table.name)


def _is_mysql_unknown_database_error(exc: Exception) -> bool:
    message = str(exc).lower()
    return "unknown database" in message or "1049" in message


async def init_engine(
    backend: str,
    *,
    url: str = "",
    echo: bool = False,
    pool_size: int = 5,
    max_overflow: int = 10,
    pool_timeout: int = 30,
    sqlite_dir: str = "",
) -> None:
    """Create the async engine and session factory, then auto-create tables.

    Args:
        backend: "memory", "sqlite", "postgres", or "mysql".
        url: SQLAlchemy async URL (for sqlite/postgres/mysql).
        echo: Echo SQL to log.
        pool_size: Base connection pool size (postgres/mysql).
        max_overflow: Extra connections allowed beyond pool_size under load (postgres/mysql).
        pool_timeout: Seconds to wait for a free pooled connection before erroring (postgres/mysql).
        sqlite_dir: Directory to create for SQLite (ensured to exist).
    """
    global _engine, _session_factory

    if backend == "memory":
        logger.info("Persistence backend=memory -- ORM engine not initialized")
        return

    if backend == "postgres":
        try:
            import asyncpg  # noqa: F401
        except ImportError:
            raise ImportError("database.backend is set to 'postgres' but asyncpg is not installed.\nInstall it with:\n    uv sync --extra postgres\nOr switch to backend: sqlite in config.yaml for single-node deployment.") from None
    if backend == "mysql":
        _add_runtime_vendor_path()
        driver = _mysql_driver_name(url)
        if find_spec(driver) is None:
            raise ImportError(
                "database.backend is set to 'mysql' but the MySQL async driver is not installed.\n"
                "Install asyncmy, or run switch_database_to_mysql.sh so it can install asyncmy into "
                f"{Path('.deer-flow/python-vendor')} from the offline wheelhouse."
            )

    if backend == "sqlite":
        import os

        from sqlalchemy import event

        os.makedirs(sqlite_dir or ".", exist_ok=True)
        _engine = create_async_engine(url, echo=echo, json_serializer=_json_serializer)

        # Enable WAL on every new connection. SQLite PRAGMA settings are
        # per-connection, so we wire the listener instead of running PRAGMA
        # once at startup. WAL gives concurrent reads + writers without
        # blocking and is the standard recommendation for any production
        # SQLite deployment (TC-UPG-06 in AUTH_TEST_PLAN.md). The companion
        # ``synchronous=NORMAL`` is the safe-and-fast pairing — fsync only
        # at WAL checkpoint boundaries instead of every commit.
        # Note: we do not set PRAGMA busy_timeout here — Python's sqlite3
        # driver already defaults to a 5-second busy timeout (see the
        # ``timeout`` kwarg of ``sqlite3.connect``), and aiosqlite /
        # SQLAlchemy's aiosqlite dialect inherit that default.  Setting
        # it again would be a no-op.
        @event.listens_for(_engine.sync_engine, "connect")
        def _enable_sqlite_wal(dbapi_conn, _record):  # noqa: ARG001 — SQLAlchemy contract
            cursor = dbapi_conn.cursor()
            try:
                cursor.execute("PRAGMA journal_mode=WAL;")
                cursor.execute("PRAGMA synchronous=NORMAL;")
                cursor.execute("PRAGMA foreign_keys=ON;")
            finally:
                cursor.close()
    elif backend == "postgres":
        _engine = create_async_engine(
            url,
            echo=echo,
            pool_size=pool_size,
            max_overflow=max_overflow,
            pool_timeout=pool_timeout,
            pool_pre_ping=True,
            json_serializer=_json_serializer,
        )
    elif backend == "mysql":
        _engine = create_async_engine(
            url,
            echo=echo,
            pool_size=pool_size,
            max_overflow=max_overflow,
            pool_timeout=pool_timeout,
            pool_pre_ping=True,
            pool_recycle=1800,
            json_serializer=_json_serializer,
        )
    else:
        raise ValueError(f"Unknown persistence backend: {backend!r}")

    _session_factory = async_sessionmaker(_engine, expire_on_commit=False)

    # Auto-create tables for fresh deployments, then let Alembic apply any
    # in-place schema fixes on top (see ``_run_alembic_upgrade`` below). Both
    # are idempotent so repeating them across restarts is harmless.
    from deerflow.persistence.base import Base

    # Import all models so Base.metadata discovers them.
    # When no models exist yet (scaffolding phase), this is a no-op.
    try:
        import deerflow.persistence.models  # noqa: F401
    except ImportError:
        # Models package not yet available — tables won't be auto-created.
        # This is expected during initial scaffolding or minimal installs.
        logger.debug("deerflow.persistence.models not found; skipping auto-create tables")

    try:
        async with _engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)
    except Exception as exc:
        if backend == "postgres" and "does not exist" in str(exc):
            # Database not yet created — attempt to auto-create it, then retry.
            await _auto_create_postgres_db(url)
            # Rebuild engine against the now-existing database
            await _engine.dispose()
            _engine = create_async_engine(url, echo=echo, pool_size=pool_size, max_overflow=max_overflow, pool_timeout=pool_timeout, pool_pre_ping=True, json_serializer=_json_serializer)
            _session_factory = async_sessionmaker(_engine, expire_on_commit=False)
            async with _engine.begin() as conn:
                await conn.run_sync(Base.metadata.create_all)
        elif backend == "mysql" and _is_mysql_unknown_database_error(exc):
            # Database not yet created. Create it, rebuild the app engine
            # against the target DB, then let SQLAlchemy create all tables.
            await _auto_create_mysql_db(url)
            await _engine.dispose()
            _engine = create_async_engine(
                url,
                echo=echo,
                pool_size=pool_size,
                max_overflow=max_overflow,
                pool_timeout=pool_timeout,
                pool_pre_ping=True,
                pool_recycle=1800,
                json_serializer=_json_serializer,
            )
            _session_factory = async_sessionmaker(_engine, expire_on_commit=False)
            async with _engine.begin() as conn:
                await conn.run_sync(Base.metadata.create_all)
        else:
            raise

    # Patch existing tables with ORM-declared columns that ``create_all``
    # cannot add (it only creates brand-new tables). Best-effort: a failure is
    # logged and startup continues on the current schema.
    try:
        async with _engine.begin() as conn:
            await conn.run_sync(_ensure_orm_columns_sync)
    except Exception:
        logger.exception("Auto-adding missing ORM columns failed; continuing with current schema")

    # Apply in-place schema fixes (column widenings, etc.) that ``create_all``
    # cannot do on its own. Failures don't abort startup — see the helper.
    # await _run_alembic_upgrade(url)

    logger.info("Persistence engine initialized: backend=%s", backend)


async def init_engine_from_config(config) -> None:
    """Convenience: init engine from a DatabaseConfig object."""
    if config.backend == "memory":
        await init_engine("memory")
        return
    await init_engine(
        backend=config.backend,
        url=config.app_sqlalchemy_url,
        echo=config.echo_sql,
        pool_size=config.pool_size,
        max_overflow=getattr(config, "max_overflow", 10),
        pool_timeout=getattr(config, "pool_timeout", 30),
        sqlite_dir=config.sqlite_dir if config.backend == "sqlite" else "",
    )


def get_session_factory() -> async_sessionmaker[AsyncSession] | None:
    """Return the async session factory, or None if backend=memory."""
    return _session_factory


def get_engine() -> AsyncEngine | None:
    """Return the async engine, or None if not initialized."""
    return _engine


async def close_engine() -> None:
    """Dispose the engine, release all connections."""
    global _engine, _session_factory
    if _engine is not None:
        await _engine.dispose()
        logger.info("Persistence engine closed")
    _engine = None
    _session_factory = None