diff --git a/backend/CLAUDE.md b/backend/CLAUDE.md index f031d7aab..142e10803 100644 --- a/backend/CLAUDE.md +++ b/backend/CLAUDE.md @@ -95,6 +95,7 @@ make test # Run all backend tests make test-blocking-io # Run strict Blockbuster runtime gate on tests/blocking_io/ make lint # Lint with ruff make format # Format code with ruff +make migrate-rev MSG="..." # Autogenerate a new alembic revision (see Schema Migrations section) ``` The `detect-blocking-io` target parses `app/`, `packages/harness/deerflow/`, @@ -487,6 +488,41 @@ Focused regression coverage for the updater lives in `backend/tests/test_memory_ - `resolve_variable(path)` - Import module and return variable (e.g., `module.path:variable_name`) - `resolve_class(path, base_class)` - Import and validate class against base class +### Schema Migrations (`packages/harness/deerflow/persistence/migrations/`) + +DeerFlow's application tables (`runs`, `threads_meta`, `feedback`, `users`, `run_events`, plus the four `channel_*` tables) are owned by alembic via a **hybrid bootstrap** strategy. LangGraph's checkpointer tables (`checkpoints`, `checkpoint_blobs`, `checkpoint_writes`, `checkpoint_migrations`) live in the same database but are owned by LangGraph and excluded from alembic's view via `migrations/_env_filters.py::include_object`. + +**Convention**: every ORM model change (new column, new table, new index) MUST ship as an alembic revision under `migrations/versions/`. The Gateway runs `alembic upgrade head` automatically on startup; users do not run `alembic` manually in production. + +**Hybrid bootstrap** (`persistence/bootstrap.py::bootstrap_schema`, invoked from `persistence/engine.py::init_engine`): + +| DB state | Action | +|-------------------------------------------|-----------------------------------------| +| empty (no DeerFlow tables) | `create_all` + `alembic stamp head` | +| legacy (DeerFlow tables, no `alembic_version`) | `create_all` (baseline tables only, backfill) + `alembic stamp 0001_baseline` + `upgrade head` | +| versioned (`alembic_version` row exists) | `alembic upgrade head` | + +The legacy branch handles pre-alembic databases that already have at least one DeerFlow-owned table. `create_all` runs first because stamping at `0001_baseline` makes alembic skip the baseline's own `create_table` DDL on the subsequent upgrade — so any baseline table introduced into `Base.metadata` after the user's DB was first provisioned (e.g. the `channel_*` tables from PR #1930 for users upgrading across multiple releases) would otherwise never be created, and the first request hitting that table would 500 with `no such table`. The backfill is **restricted to `_BASELINE_TABLE_NAMES`** so it does not also create tables that future revisions introduce — those revisions' own `op.create_table` would otherwise fail with `relation already exists`. A guard test pins `_BASELINE_TABLE_NAMES` against `0001_baseline.upgrade()`'s actual output, so editing 0001 to add or remove a table forces a matching update to the constant. Column-level shape (pre-#3658 vs post-#3658 vs manual-ALTER for `token_usage_by_model`) is answered by each `versions/*.py` revision via the idempotent helpers in `migrations/_helpers.py` (`safe_add_column` / `safe_drop_column`) which no-op when the change is already present and `logger.warning` on shape drift. **Adding a new ORM column / table only requires a new revision file — no edit to `bootstrap.py` is needed** *unless* the new revision adds a new baseline table (rare; only happens when a new model is part of the baseline rather than introduced by its own revision). + +The empty-DB path keeps using `create_all` because `Base.metadata` is the only authoritative schema source — `create_all` renders both SQLite (JSON, type affinity) and Postgres (JSONB, partial indexes) correctly without anyone having to keep a hand-written baseline in lockstep. `0001_baseline.upgrade()` is therefore almost never executed in practice; it exists as a stamp target + chain root. + +**Concurrency safety**: Postgres uses `pg_advisory_lock` to serialise concurrent Gateway instances. SQLite uses a per-engine `asyncio.Lock` for same-process startup and is best-effort across processes via SQLite's file-level write lock + `PRAGMA busy_timeout`; multi-instance deployments should use Postgres. Column revisions in `versions/` additionally use idempotent helpers (`_helpers.py::safe_add_column`, `safe_drop_column`) so repeated post-baseline changes and retries are no-ops when the change is already present. + +**Authoring a new revision**: +```bash +cd backend && make migrate-rev MSG="add foo column to runs" +``` +This invokes `alembic revision --autogenerate` against the live ORM models. Review the generated file under `migrations/versions/` and switch raw `op.add_column` / `op.drop_column` calls to the idempotent helpers from `_helpers.py` before committing. There is no `make migrate` / `make migrate-stamp` target on purpose — the only execution path is Gateway startup, which keeps operational mistakes off the table. + +**Where things live**: +- `migrations/env.py` — alembic env, delegates filter to `_env_filters.py`, sets `render_as_batch=True` for SQLite ALTER support +- `migrations/_env_filters.py::include_object` — drops LangGraph checkpointer tables from alembic's view +- `migrations/_helpers.py` — `safe_add_column` / `safe_drop_column` +- `migrations/versions/0001_baseline.py` — chain root, matches the schema `create_all` produces from `Base.metadata` +- `migrations/versions/0002_runs_token_usage.py` — fixes issue #3682 +- `persistence/bootstrap.py` — `bootstrap_schema(engine, backend=...)`, the three-branch decision + locking +- Tests: `tests/test_persistence_bootstrap.py` (branches), `tests/test_persistence_bootstrap_concurrency.py` (concurrency), `tests/test_persistence_bootstrap_regression.py` (issue #3682), `tests/test_persistence_migrations_env.py` (filter), `tests/blocking_io/test_persistence_bootstrap.py` (asyncio.to_thread anchor) + ### Tracing System (`packages/harness/deerflow/tracing/`) LangSmith and Langfuse are both supported. The wiring lives in two layers: diff --git a/backend/Makefile b/backend/Makefile index a8ecc602c..f903f9e5b 100644 --- a/backend/Makefile +++ b/backend/Makefile @@ -22,3 +22,14 @@ format: detect-blocking-io: @PYTHONPATH=. PYTHONIOENCODING=utf-8 PYTHONUTF8=1 uv run python ../scripts/detect_blocking_io_static.py --output ../.deer-flow/blocking-io-findings.json + +# Generate a new alembic revision by autogenerating the diff against the live +# ORM models. Usage: make migrate-rev MSG="add foo column to runs" +# The Gateway runs `alembic upgrade head` automatically at startup via +# `bootstrap_schema`, so there is intentionally no `migrate` / `migrate-stamp` +# target -- the single execution path keeps ops mistakes off the table. +# The script builds a fresh temp SQLite at head, then diffs the live models +# against it, so a clean checkout does NOT need a pre-existing ./data/deerflow.db. +migrate-rev: + @if [ -z "$(MSG)" ]; then echo 'usage: make migrate-rev MSG="describe the change"'; exit 1; fi + PYTHONPATH=. PYTHONIOENCODING=utf-8 PYTHONUTF8=1 uv run python scripts/_autogen_revision.py "$(MSG)" diff --git a/backend/README.md b/backend/README.md index 22eb8f71d..cb48b93a0 100644 --- a/backend/README.md +++ b/backend/README.md @@ -364,8 +364,35 @@ make gateway # Run Gateway API without reload (port 8001) make lint # Run linter (ruff) make format # Format code (ruff) make detect-blocking-io # Inventory blocking IO that may block the backend event loop +make migrate-rev MSG="..." # Autogenerate a new alembic revision against the live ORM models ``` +### Schema Migrations + +DeerFlow's application tables (`runs`, `threads_meta`, `feedback`, `users`, +`run_events`, and the `channel_*` tables) are owned by alembic. The Gateway +runs `alembic upgrade head` automatically on startup via +`bootstrap_schema(engine, backend=...)`, so operators do not run `alembic` +manually in production. Bootstrap is concurrency-safe (Postgres advisory lock +across processes; per-engine `asyncio.Lock` inside one SQLite process) and +idempotent against pre-existing schemas (empty / legacy / versioned). + +When you add or change an ORM model, ship the change as a new revision under +`packages/harness/deerflow/persistence/migrations/versions/`: + +```bash +make migrate-rev MSG="add foo column to runs" +``` + +The target invokes `scripts/_autogen_revision.py`, which builds a fresh temp +SQLite at `head` and diffs the live models against it — so a clean checkout +does not need a pre-existing `./data/deerflow.db`. Review the generated file +and switch raw `op.add_column` / `op.drop_column` calls to the idempotent +helpers in `migrations/_helpers.py` before committing. There is no +`make migrate` / `make migrate-stamp` target on purpose — Gateway startup is +the only execution path, which keeps operational mistakes off the table. See +`backend/CLAUDE.md` (Schema Migrations) for the full design. + ### Code Style - **Linter/Formatter**: `ruff` diff --git a/backend/packages/harness/deerflow/persistence/bootstrap.py b/backend/packages/harness/deerflow/persistence/bootstrap.py new file mode 100644 index 000000000..c4f530ae7 --- /dev/null +++ b/backend/packages/harness/deerflow/persistence/bootstrap.py @@ -0,0 +1,452 @@ +"""Hybrid schema bootstrap for DeerFlow's application tables. + +Replaces the unconditional ``Base.metadata.create_all`` at Gateway startup. +Combines two ideas: + +1. ``create_all`` stays the empty-DB fast path -- it renders ``Base.metadata`` + faithfully across SQLite and Postgres dialects (JSON vs JSONB, server + defaults, index/FK names, type affinity) without anyone having to hand-keep + a mirror baseline in sync with the models. +2. **Alembic owns every change from baseline onward.** Any new ORM column / + table / index must ship as a revision under ``migrations/versions/``. + +Three-branch decision (see ``_decide_state``) +--------------------------------------------- + +| DB state | Action | +|---------------------------------------|-----------------------------------------| +| empty (no DeerFlow tables) | ``create_all`` + ``alembic stamp head`` | +| legacy (DeerFlow tables, no alembic) | ``create_all`` (baseline tables only, as backfill) + ``stamp 0001_baseline`` + ``upgrade head`` | +| versioned (``alembic_version`` row) | ``alembic upgrade head`` | + +The legacy branch handles pre-alembic databases that already have at least one +DeerFlow-owned table. ``create_all`` runs first because stamping at +``0001_baseline`` makes alembic skip the baseline's own ``create_table`` DDL on +the subsequent upgrade -- so any baseline table introduced into +``Base.metadata`` after the user's DB was first provisioned (e.g. the +``channel_*`` tables from PR #1930 for users upgrading across multiple +releases) would otherwise never be created, and the first request hitting that +table would 500 with ``no such table``. The backfill is **restricted to +``_BASELINE_TABLE_NAMES``** so it does not also create tables that future +revisions introduce -- those revisions' own ``op.create_table`` would then +fail with ``relation already exists``. A guard test pins the restriction +set against ``0001_baseline.upgrade()``'s actual output. + +Column-level shape (the pre-#3658 vs post-#3658 vs manual-ALTER cases for +``token_usage_by_model``) is answered by each ``versions/*.py`` revision via +the idempotent helpers in ``migrations/_helpers.py`` (``safe_add_column`` +no-ops when the column is already present and ``logger.warning``s on +shape drift). Future schema additions therefore plug in by writing a new +revision file -- **no edit to this module is required** *unless* the new +revision creates a new baseline table, in which case ``_BASELINE_TABLE_NAMES`` +must be updated to match (the guard test fires otherwise). + +Concurrency safety +------------------ + +Layered, with different guarantees per backend. Postgres has true +cross-process serialisation. SQLite is single-process safe and cross-process +best-effort; multi-instance deployments should use Postgres. + +* **Postgres -- true cross-process serialisation.** ``pg_advisory_lock`` runs + the whole reflect-and-act sequence under an exclusive lock that survives + cross-process. Concurrent Gateway instances queue cleanly and the second + one observes head as a no-op. + +* **SQLite -- single-process serialisation, best-effort cross-process.** + SQLite is single-node by deployment, so the realistic concurrency case is + multiple async tasks inside one Gateway process (tests, lifespan re-entry). + A per-engine ``asyncio.Lock`` serialises those. For the rare cross-process + case (e.g. two ``make dev`` workers on the same DB file), we rely on + SQLite's own file-level write lock plus a 30s ``PRAGMA busy_timeout`` -- + the latter is set on **both** the production engine + (``persistence/engine.py``) and the alembic-spawned engine + (``migrations/env.py``) so any writer waits up to 30s for the file lock + instead of failing fast. This is best-effort, not a true mutex: under + pathological overlap a process can still see ``database is locked`` after + 30s. The fallback line of defence -- idempotent revisions -- guarantees + correctness anyway. + +* **Idempotent revisions -- retry fallback.** Column revisions use the helpers + in ``migrations/_helpers.py`` so repeated post-baseline changes, manual + ALTERs, or retries after SQLite lock contention do not duplicate work. + +``alembic upgrade head`` on a DB already at head is a no-op by alembic's own +semantics, so the second-N-th actor simply observes head and exits. +""" + +from __future__ import annotations + +import asyncio +import logging +import weakref +from contextlib import asynccontextmanager +from pathlib import Path +from typing import Any + +from alembic import command as alembic_command +from alembic.config import Config as AlembicConfig +from alembic.script import ScriptDirectory +from sqlalchemy import inspect as sa_inspect +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncEngine + +logger = logging.getLogger(__name__) + + +# Where the alembic environment lives, relative to this file. +_MIGRATIONS_DIR = Path(__file__).resolve().parent / "migrations" + +# Cached migration head, computed once per process from the disk script tree. +_HEAD_REVISION: str | None = None + +# Baseline (stamp target for legacy DBs). Pinned here so the bootstrap layer +# fails loudly if the baseline revision is ever renamed without updating the +# stamp call. ``tests/test_persistence_bootstrap.py`` asserts this string is a +# real revision id in the script tree. +_BASELINE_REVISION = "0001_baseline" + +# Stable advisory-lock key for Postgres. Two random 32-bit halves picked once +# so we never collide with any other application's advisory locks. Do not +# change without coordinating a one-time migration (a key change effectively +# releases the prior lock). +_PG_LOCK_KEY = 0x0DEE_12F1_0BEE_3682 + + +# Tables created by ``0001_baseline.upgrade()``. The legacy branch restricts +# its ``create_all`` backfill to this set so it does NOT pre-empt later +# ``op.create_table`` revisions for models added after baseline -- those +# revisions would otherwise fail with ``relation already exists`` if +# ``create_all`` had created their table first. (Column revisions are +# already safe via the idempotent helpers in ``migrations/_helpers.py``; +# there is no analogous ``safe_create_table`` yet, so we keep table-level +# safety at this layer instead of pushing it onto every future revision.) +# +# ``test_baseline_table_names_constant_matches_0001`` pins this set against +# what 0001 actually creates -- editing 0001 without updating this constant +# (or vice versa) fires that test. +_BASELINE_TABLE_NAMES: frozenset[str] = frozenset( + { + "channel_connections", + "channel_conversations", + "channel_credentials", + "channel_oauth_states", + "feedback", + "run_events", + "runs", + "threads_meta", + "users", + } +) + + +# Per-engine SQLite bootstrap locks. Per-engine (not module-global) so each +# engine instance pairs with a lock bound to the event loop that uses that +# engine -- necessary because ``asyncio.Lock`` binds to the first loop it sees, +# and pytest gives each async test its own loop. Production uses one engine +# per process so this dict collapses to a single entry in practice. +# +# Keyed by the engine object itself via ``WeakKeyDictionary`` rather than +# ``id(engine)``: CPython recycles addresses after GC, so a stale ``id`` → +# ``Lock`` entry from a dead engine could be returned to a new engine that +# happened to land on the same address. The returned lock would still be bound +# to the dead engine's event loop and ``async with`` would raise +# ``RuntimeError: ... bound to a different event loop``. Hashing the engine +# itself also drops entries automatically when the engine is collected, so this +# dict never grows past the live engine count. +_SQLITE_LOCKS: weakref.WeakKeyDictionary[AsyncEngine, asyncio.Lock] = weakref.WeakKeyDictionary() + + +def _get_sqlite_local_lock(engine: AsyncEngine) -> asyncio.Lock: + lock = _SQLITE_LOCKS.get(engine) + if lock is None: + lock = asyncio.Lock() + _SQLITE_LOCKS[engine] = lock + return lock + + +def _escape_url_for_alembic(url: str) -> str: + """Double literal ``%`` so ``ConfigParser`` interpolation leaves the URL intact. + + ``alembic.config.Config.set_main_option`` forwards to ``ConfigParser.set``, + which performs ``%(name)s``-style interpolation on the value. A URL-encoded + password like ``p%40ss`` (``@`` escaped to ``%40``) would otherwise raise + ``InterpolationSyntaxError``. Doubling every literal ``%`` makes + ConfigParser unescape it back to one. Shared with + ``scripts/_autogen_revision.py`` so the round-trip rule lives in one place. + """ + return url.replace("%", "%%") + + +def _alembic_safe_url(engine: AsyncEngine) -> str: + """Render *engine*'s URL in a form alembic ``set_main_option`` accepts. + + Two pitfalls handled: + + 1. ``str(engine.url)`` (and ``URL.render_as_string()`` without args) masks + the password as ``***`` -- so alembic's stamp/upgrade would open its own + connection with garbage credentials and fail at runtime, even though + the live engine connects fine. Fix: ``render_as_string(hide_password=False)``. + 2. ConfigParser interpolation on ``%`` -- delegated to + ``_escape_url_for_alembic`` so the rule is shared with the autogen + script. + """ + rendered = engine.url.render_as_string(hide_password=False) + return _escape_url_for_alembic(rendered) + + +def _get_alembic_config(engine: AsyncEngine) -> AlembicConfig: + """Build an in-process alembic config pointing at our migrations dir. + + Avoids reading ``alembic.ini`` from disk so the production runtime doesn't + depend on a working-directory-relative file lookup. The ``script_location`` + is anchored at the package path on disk. + """ + cfg = AlembicConfig() + cfg.set_main_option("script_location", str(_MIGRATIONS_DIR)) + cfg.set_main_option("sqlalchemy.url", _alembic_safe_url(engine)) + return cfg + + +def _get_head_revision() -> str: + """Return the head revision id from ``versions/``, cached per process.""" + global _HEAD_REVISION + if _HEAD_REVISION is None: + cfg = AlembicConfig() + cfg.set_main_option("script_location", str(_MIGRATIONS_DIR)) + script = ScriptDirectory.from_config(cfg) + head = script.get_current_head() + if head is None: + raise RuntimeError("alembic has no head revision -- versions/ directory is empty") + _HEAD_REVISION = head + return _HEAD_REVISION + + +def _reflect_state(sync_conn: Any) -> dict[str, bool]: + """Inspect *sync_conn* (sync connection inside ``run_sync``) and return: + + - ``has_alembic_version``: bool + - ``has_deerflow_tables``: True iff at least one table that ``Base.metadata`` + knows about is present in the DB. Computed as ``reflected ∩ metadata`` so + the bootstrap layer never hardcodes a specific table or column name -- + adding a new ORM model only changes ``Base.metadata``, not this module. + """ + from deerflow.persistence.base import Base + + # Make sure every ORM model is imported, otherwise ``Base.metadata.tables`` + # may miss tables registered by submodules that haven't been imported yet. + try: + import deerflow.persistence.models # noqa: F401 + except ImportError: + logger.debug("deerflow.persistence.models not found; metadata may be incomplete") + + insp = sa_inspect(sync_conn) + reflected = set(insp.get_table_names()) + metadata_tables = set(Base.metadata.tables) + return { + "has_alembic_version": "alembic_version" in reflected, + "has_deerflow_tables": bool(reflected & metadata_tables), + } + + +def _decide_state(state: dict[str, bool]) -> str: + """Map a reflected DB state to one of three branch labels. + + The legacy branch covers every pre-alembic DB uniformly -- whether the + columns added by later revisions are present or not is a question each + revision answers for itself via the idempotent helpers in + ``migrations/_helpers.py``. + """ + if state["has_alembic_version"]: + return "versioned" + if not state["has_deerflow_tables"]: + # Either a brand-new DB or a DB containing only tables we don't own + # (e.g. LangGraph's checkpointer tables on a fresh deployment). The + # empty branch provisions the tables alembic owns, then stamps head. + return "empty" + return "legacy" + + +def _run_create_all_sync(sync_conn: Any) -> None: + """Create all DeerFlow-owned tables on *sync_conn*.""" + # Import here to ensure all model classes are registered with Base.metadata. + from deerflow.persistence.base import Base + + try: + import deerflow.persistence.models # noqa: F401 + except ImportError: + logger.debug("deerflow.persistence.models not found; bootstrap will create empty schema") + + Base.metadata.create_all(sync_conn) + + +def _run_baseline_create_all_sync(sync_conn: Any) -> None: + """Create only the baseline tables on *sync_conn* (idempotent via checkfirst). + + Used by the legacy branch to backfill baseline-era tables missing from + the user's DB. Restricting the table list to ``_BASELINE_TABLE_NAMES`` + is the safety property: an unrestricted ``create_all`` would also create + tables introduced by later revisions, which would then collide with + those revisions' ``op.create_table`` calls when alembic ran upgrade. + """ + from deerflow.persistence.base import Base + + try: + import deerflow.persistence.models # noqa: F401 + except ImportError: + logger.debug("deerflow.persistence.models not found; baseline backfill may be incomplete") + + baseline_tables = [Base.metadata.tables[name] for name in _BASELINE_TABLE_NAMES if name in Base.metadata.tables] + Base.metadata.create_all(sync_conn, tables=baseline_tables, checkfirst=True) + + +def _stamp(cfg: AlembicConfig, revision: str) -> None: + """Synchronous alembic stamp; callers must wrap in ``asyncio.to_thread``.""" + alembic_command.stamp(cfg, revision) + + +def _upgrade(cfg: AlembicConfig, revision: str) -> None: + """Synchronous alembic upgrade; callers must wrap in ``asyncio.to_thread``.""" + alembic_command.upgrade(cfg, revision) + + +# --------------------------------------------------------------------------- +# Cross-process locking +# --------------------------------------------------------------------------- + + +@asynccontextmanager +async def _postgres_lock(engine: AsyncEngine): + """Hold a Postgres session-level advisory lock for the body of the block. + + Session-level (not transaction-level) so the lock outlives implicit + transactions opened by alembic during ``stamp`` / ``upgrade``. The lock + is released explicitly on the way out and -- as a safety net -- when the + backing session disconnects (process crash, kill -9). + + Idle-in-transaction protection + ------------------------------ + + ``engine.connect()`` auto-begins a transaction on the first ``execute``, + and this connection then sits idle while ``asyncio.to_thread(_upgrade, + ...)`` runs alembic on a *different* pooled connection. Managed Postgres + (RDS, Cloud SQL, Supabase) ships with ``idle_in_transaction_session_ + timeout`` set to 1-10 minutes by default; if alembic takes longer than + that, the host kills this idle-in-transaction session, and because + advisory locks are session-scoped, the lock is **silently released**. + A second Gateway then acquires it and runs DDL concurrently with the + first -- defeating the whole purpose of the lock. + + Defence: ``SET LOCAL idle_in_transaction_session_timeout = 0`` disables + the kill **for this transaction only** (no global / role-level effect). + Self-hosted Postgres usually ships with the timeout off, so this is a + no-op there; on managed PG it is what keeps the lock alive while DDL + runs. Must execute *before* ``pg_advisory_lock`` so a slow lock acquire + on a heavily-contended cluster is itself protected. + """ + async with engine.connect() as conn: + await conn.execute(text("SET LOCAL idle_in_transaction_session_timeout = 0")) + await conn.execute(text("SELECT pg_advisory_lock(:k)"), {"k": _PG_LOCK_KEY}) + try: + logger.info("bootstrap: acquired postgres advisory lock key=0x%x", _PG_LOCK_KEY) + yield + finally: + try: + await conn.execute(text("SELECT pg_advisory_unlock(:k)"), {"k": _PG_LOCK_KEY}) + except Exception: # noqa: BLE001 + logger.warning("bootstrap: pg_advisory_unlock raised; session close will release", exc_info=True) + + +@asynccontextmanager +async def _sqlite_lock(engine: AsyncEngine): + """Serialise SQLite bootstrap inside one process; cross-process is + best-effort via SQLite's own file lock + ``PRAGMA busy_timeout``. + + Why not ``BEGIN IMMEDIATE`` on a sentinel connection? SQLite is + single-writer per file. If we held a write lock on one connection, + alembic's own connection (opened inside ``stamp`` / ``upgrade``) would + deadlock against us. + + Why not a cross-process OS file lock? It would work, but it adds a hard + dependency on platform-specific ``fcntl`` / ``msvcrt`` calls for a + deployment shape (multi-process SQLite) that's already discouraged for + DeerFlow. The 30s ``busy_timeout`` plus idempotent revisions cover the + realistic case; truly multi-instance deployments should use Postgres. + + Note: the 30s ``busy_timeout`` is set by the engine event hooks in + ``persistence/engine.py`` (production) and ``migrations/env.py`` + (alembic-spawned). This function relies on those PRAGMAs being in place + rather than setting one on a probe connection that wouldn't propagate. + """ + async with _get_sqlite_local_lock(engine): + logger.info("bootstrap: acquired sqlite in-process lock") + yield + + +def _bootstrap_lock(engine: AsyncEngine, *, backend: str): + if backend == "postgres": + return _postgres_lock(engine) + if backend == "sqlite": + return _sqlite_lock(engine) + raise ValueError(f"bootstrap: unsupported backend {backend!r}") + + +# --------------------------------------------------------------------------- +# Top-level entry point +# --------------------------------------------------------------------------- + + +async def bootstrap_schema(engine: AsyncEngine, *, backend: str) -> None: + """Bring the DB schema to head. + + Postgres calls are serialised across processes with an advisory lock. + SQLite calls are serialised inside one process and are best-effort across + processes via SQLite's file lock and ``busy_timeout``. + + Branch dispatch is documented at module top. ``alembic.command.stamp`` and + ``alembic.command.upgrade`` are synchronous and would block the event + loop; both are wrapped in ``asyncio.to_thread``. + """ + head = _get_head_revision() + cfg = _get_alembic_config(engine) + + async with _bootstrap_lock(engine, backend=backend): + async with engine.connect() as conn: + state = await conn.run_sync(_reflect_state) + decision = _decide_state(state) + + if decision == "empty": + logger.info("bootstrap: branch=empty -> create_all + stamp head (%s)", head) + async with engine.begin() as conn: + await conn.run_sync(_run_create_all_sync) + await asyncio.to_thread(_stamp, cfg, head) + + elif decision == "legacy": + logger.info( + "bootstrap: branch=legacy -> create_all (backfill missing baseline tables) + stamp %s + upgrade head (%s)", + _BASELINE_REVISION, + head, + ) + # ``_run_baseline_create_all_sync`` is restricted to + # ``_BASELINE_TABLE_NAMES`` -- a plain ``Base.metadata.create_all`` + # would also create tables introduced by later revisions and + # collide with their ``op.create_table`` on the subsequent + # upgrade. With the restriction, missing baseline tables are + # backfilled and post-baseline ``create_table`` revisions run + # against a DB where their tables genuinely do not yet exist. + # The post-create_all column-add revisions still no-op via + # ``safe_add_column`` because baseline-era tables now have the + # columns those revisions would add. + async with engine.begin() as conn: + await conn.run_sync(_run_baseline_create_all_sync) + await asyncio.to_thread(_stamp, cfg, _BASELINE_REVISION) + await asyncio.to_thread(_upgrade, cfg, "head") + + elif decision == "versioned": + logger.info("bootstrap: branch=versioned -> upgrade head (%s)", head) + await asyncio.to_thread(_upgrade, cfg, "head") + + else: # pragma: no cover -- defensive + raise RuntimeError(f"bootstrap: unhandled decision {decision!r}") + + logger.info("bootstrap: complete (backend=%s)", backend) diff --git a/backend/packages/harness/deerflow/persistence/engine.py b/backend/packages/harness/deerflow/persistence/engine.py index 9aa58d7f3..0e12f57ef 100644 --- a/backend/packages/harness/deerflow/persistence/engine.py +++ b/backend/packages/harness/deerflow/persistence/engine.py @@ -112,11 +112,14 @@ async def init_engine( # 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. + # We also widen ``busy_timeout`` to 30s here. Python's sqlite3 driver + # defaults to 5s, which is fine for transient row contention but too + # tight for cross-process bootstrap: the second-N-th Gateway process + # may need to wait while the first runs ``ALTER TABLE`` / + # ``CREATE TABLE`` for a fresh schema. The same widened timeout is + # mirrored on the alembic-spawned engine in + # ``migrations/env.py::run_migrations_online`` so its connections + # behave identically. @event.listens_for(_engine.sync_engine, "connect") def _enable_sqlite_wal(dbapi_conn, _record): # noqa: ARG001 — SQLAlchemy contract cursor = dbapi_conn.cursor() @@ -124,6 +127,7 @@ async def init_engine( cursor.execute("PRAGMA journal_mode=WAL;") cursor.execute("PRAGMA synchronous=NORMAL;") cursor.execute("PRAGMA foreign_keys=ON;") + cursor.execute("PRAGMA busy_timeout=30000;") finally: cursor.close() elif backend == "postgres": @@ -139,31 +143,28 @@ async def init_engine( _session_factory = async_sessionmaker(_engine, expire_on_commit=False) - # Auto-create tables (dev convenience). Production should use Alembic. - 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") + # Schema bootstrap (hybrid): + # - empty DB -> create_all + alembic stamp head + # - legacy DB -> create_all (baseline tables only, backfill) + alembic stamp baseline + upgrade head + # - already managed -> alembic upgrade head + # Concurrency: Postgres advisory lock (true cross-process); SQLite uses an + # in-process asyncio.Lock plus a 30s PRAGMA busy_timeout (also set on + # alembic's own connections in env.py) -- multi-process SQLite bootstrap + # is best-effort, gated by SQLite's natural file-level write lock. + # See deerflow.persistence.bootstrap for the full state machine. + from deerflow.persistence.bootstrap import bootstrap_schema try: - async with _engine.begin() as conn: - await conn.run_sync(Base.metadata.create_all) + await bootstrap_schema(_engine, backend=backend) 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. + # 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, 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) + await bootstrap_schema(_engine, backend=backend) else: raise diff --git a/backend/packages/harness/deerflow/persistence/migrations/_env_filters.py b/backend/packages/harness/deerflow/persistence/migrations/_env_filters.py new file mode 100644 index 000000000..35b7f398f --- /dev/null +++ b/backend/packages/harness/deerflow/persistence/migrations/_env_filters.py @@ -0,0 +1,36 @@ +"""Object filters used by ``env.py`` to scope alembic to DeerFlow tables. + +LangGraph checkpointer tables live in the same database but are owned by +LangGraph. Without this filter, ``alembic revision --autogenerate`` would +reflect them and emit spurious ``drop_table`` ops every revision. + +Kept in its own module (instead of inlined in ``env.py``) so it can be +unit-tested without dragging in alembic's import-time machinery. +""" + +from __future__ import annotations + +# Tables owned by LangGraph -- alembic must never propose DDL for them. +LANGGRAPH_OWNED_TABLES: frozenset[str] = frozenset( + { + "checkpoints", + "checkpoint_blobs", + "checkpoint_writes", + "checkpoint_migrations", + } +) + + +def include_object(object_, name, type_, reflected, compare_to): # noqa: ARG001 + """Returns False for any LangGraph-owned table or for an index/constraint + whose parent table is LangGraph-owned. Returns True otherwise. + + Signature matches alembic's ``include_object`` callable contract: + ``(object, name, type_, reflected, compare_to)``. + """ + if type_ == "table" and name in LANGGRAPH_OWNED_TABLES: + return False + parent_table = getattr(object_, "table", None) + if parent_table is not None and getattr(parent_table, "name", None) in LANGGRAPH_OWNED_TABLES: + return False + return True diff --git a/backend/packages/harness/deerflow/persistence/migrations/_helpers.py b/backend/packages/harness/deerflow/persistence/migrations/_helpers.py new file mode 100644 index 000000000..9428a8a8e --- /dev/null +++ b/backend/packages/harness/deerflow/persistence/migrations/_helpers.py @@ -0,0 +1,189 @@ +"""Idempotent helpers for alembic column revisions. + +Column revisions in ``versions/`` should use these helpers instead of raw +``op.add_column`` / ``op.drop_column`` so re-running a column change against a +DB that already has (or has already removed) the column is a safe no-op. + +Two reasons we need idempotency: + +1. **Defence-in-depth on top of bootstrap locking.** ``bootstrap_schema()`` + serialises Postgres with an advisory lock and SQLite within one process + with an ``asyncio.Lock``. If a retry happens anyway (manual ALTER, + misconfiguration, SQLite cross-process contention), the revision must still + be safe to re-run. + +2. **Same posture that made ``Base.metadata.create_all`` forgiving.** + ``create_all`` skips existing tables. Column migrations should mirror that + forgiving behavior by skipping columns already in the desired state. + +Drift warning +------------- + +Name-match alone can hide a column that a manual ``ALTER`` (for example the +#3682 workaround that ran ``ALTER TABLE runs ADD COLUMN token_usage_by_model +JSON`` without ``NOT NULL DEFAULT '{}'``, or the wrong-type variant +``ALTER TABLE runs ADD COLUMN token_usage_by_model TEXT NOT NULL DEFAULT +'{}'``) left in a shape that diverges from what ``Base.metadata.create_all`` +would produce on a fresh DB. To surface that silent drift, ``safe_add_column`` +compares the existing column's ``nullable`` / ``server_default`` / ``type`` +against the desired ``sa.Column`` and emits ``logger.warning`` on mismatch. +Type comparison goes through ``_type_equivalent``, which treats known +dialect-synonym pairs (e.g. ``JSON`` vs ``JSONB``) as equivalent to avoid +false positives while still catching wholesale type mismatches like +``TEXT`` vs ``JSON``. We do not auto-repair -- a warning is enough for +operators to notice and decide. +""" + +from __future__ import annotations + +import logging + +import sqlalchemy as sa +from alembic import op + +logger = logging.getLogger(__name__) + + +def _inspector() -> sa.Inspector: + return sa.inspect(op.get_bind()) + + +def _normalize_default(value: object) -> str | None: + """Normalize a server-default value for cross-source comparison. + + The desired value comes from ``sa.Column.server_default`` (a + ``DefaultClause`` / ``TextClause`` literal, ``None``, or a Python literal); + the reflected value comes from ``Inspector.get_columns()['default']`` as a + dialect-rendered string. Strip outer parens / whitespace / Postgres-style + type casts so textually-equivalent forms compare equal across dialects. + """ + if value is None: + return None + if isinstance(value, sa.sql.elements.TextClause): + text = value.text + elif isinstance(value, sa.schema.DefaultClause) and isinstance(value.arg, sa.sql.elements.TextClause): + text = value.arg.text + else: + text = str(value) + text = text.strip() + # Strip a single layer of outer parens that some dialects wrap defaults in. + if text.startswith("(") and text.endswith(")"): + text = text[1:-1].strip() + # Strip Postgres-style type casts like ``'{}'::jsonb``. + if "::" in text: + text = text.split("::", 1)[0].strip() + return text or None + + +def _normalize_type(value: object) -> str: + """Normalize a SQLAlchemy ``TypeEngine`` (or reflected type) for comparison. + + Returns the upper-cased type-class name with any parameters stripped + (e.g. ``JSON()`` → ``"JSON"``, ``VARCHAR(255)`` → ``"VARCHAR"``). Length + parameters are dropped on purpose: drift warnings target wholesale type + misconfigurations (the JSON-vs-TEXT review case), not dialect-rendered + size defaults. An empty string signals "missing info" -- callers should + not equality-check empty strings. + """ + if value is None: + return "" + s = value if isinstance(value, str) else repr(value) + return s.upper().split("(", 1)[0].strip() + + +# Known dialect-synonym pairs that must NOT fire a type-drift warning. +# Postgres reflects ``JSON`` as ``JSONB`` (and vice versa depending on how +# the column was provisioned); the model's ``sa.JSON`` plus this allowlist +# keeps a Postgres deployment quiet while still catching genuine type errors +# like ``TEXT NOT NULL DEFAULT '{}'`` re-adds. +# +# Add a new pair here ONLY when a real reflection-vs-model mismatch is +# proven to be a false positive in a deployment -- not pre-emptively, since +# overly broad equivalence would re-open the silent-drift hole this helper +# exists to close. +_EQUIVALENT_TYPE_FAMILIES: tuple[frozenset[str], ...] = (frozenset({"JSON", "JSONB"}),) + + +def _type_equivalent(actual: object, desired: object) -> bool: + """True if *actual* and *desired* are the same type or a known equivalent. + + Returns True when either side is missing reflected info so missing-data + cases never false-positive into a noisy warning. + """ + a = _normalize_type(actual) + d = _normalize_type(desired) + if not a or not d: + return True + if a == d: + return True + pair = frozenset({a, d}) + return any(pair <= fam for fam in _EQUIVALENT_TYPE_FAMILIES) + + +def _check_column_drift(table: str, desired: sa.Column, actual: dict) -> None: + """Warn if an existing column's attributes diverge from the desired model. + + Equality is checked on ``nullable`` and ``server_default`` directly, and + on ``type`` via ``_type_equivalent`` (which treats known dialect-synonym + pairs like ``JSON`` vs ``JSONB`` as equivalent). The reflected and + desired type reprs are also echoed in the warning payload regardless of + whether type was the failing dimension, so an operator triaging the log + line sees the type context at a glance. + """ + diffs: list[str] = [] + + desired_nullable = True if desired.nullable is None else bool(desired.nullable) + actual_nullable = bool(actual.get("nullable", True)) + if desired_nullable != actual_nullable: + diffs.append(f"nullable actual={actual_nullable} desired={desired_nullable}") + + desired_default = _normalize_default(desired.server_default) + actual_default = _normalize_default(actual.get("default")) + if desired_default != actual_default: + diffs.append(f"server_default actual={actual_default!r} desired={desired_default!r}") + + if not _type_equivalent(actual.get("type"), desired.type): + diffs.append(f"type actual={_normalize_type(actual.get('type'))!r} desired={_normalize_type(desired.type)!r}") + + if diffs: + logger.warning( + "safe_add_column: %s.%s already exists but drifts from the model definition (%s); actual_type=%r desired_type=%r; leaving as-is -- a manual ALTER may be needed to match the model.", + table, + desired.name, + "; ".join(diffs), + actual.get("type"), + desired.type, + ) + + +def safe_add_column(table: str, column: sa.Column) -> None: + """``op.add_column`` that no-ops when the table or column is missing/present. + + - Missing table => nothing to add to. Skip silently because bootstrap only + supports legacy DBs that already have the baseline table set. + - Column already exists => no-op. Before returning, ``_check_column_drift`` + compares the existing column's nullability / server_default / type + against the desired ``column`` and ``logger.warning``\\ s on mismatch so + manually-applied workarounds do not silently survive as latent drift. + """ + insp = _inspector() + if table not in insp.get_table_names(): + return + existing = {c["name"]: c for c in insp.get_columns(table)} + if column.name in existing: + _check_column_drift(table, column, existing[column.name]) + return + with op.batch_alter_table(table) as batch: + batch.add_column(column) + + +def safe_drop_column(table: str, column_name: str) -> None: + """``op.drop_column`` that no-ops when the table or column is already gone.""" + insp = _inspector() + if table not in insp.get_table_names(): + return + existing = {c["name"] for c in insp.get_columns(table)} + if column_name not in existing: + return + with op.batch_alter_table(table) as batch: + batch.drop_column(column_name) diff --git a/backend/packages/harness/deerflow/persistence/migrations/env.py b/backend/packages/harness/deerflow/persistence/migrations/env.py index 22d053ee7..13b436b12 100644 --- a/backend/packages/harness/deerflow/persistence/migrations/env.py +++ b/backend/packages/harness/deerflow/persistence/migrations/env.py @@ -1,8 +1,15 @@ """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. +ONLY manages DeerFlow's tables (runs, threads_meta, feedback, users, +run_events, channel_connections, channel_credentials, channel_oauth_states, +channel_conversations). + +LangGraph's checkpointer tables (``checkpoints``, ``checkpoint_blobs``, +``checkpoint_writes``, ``checkpoint_migrations``) are managed by LangGraph +itself -- they have their own schema lifecycle and must not be touched by +Alembic. The ``include_object`` filter below explicitly excludes them so a +future ``alembic revision --autogenerate`` will not emit ``drop_table`` for +tables it does not own. """ from __future__ import annotations @@ -15,6 +22,14 @@ from alembic import context from sqlalchemy.ext.asyncio import create_async_engine from deerflow.persistence.base import Base +from deerflow.persistence.migrations._env_filters import ( + LANGGRAPH_OWNED_TABLES, + include_object, +) + +# Re-export under the module namespace for any consumer that addresses them +# via ``env.LANGGRAPH_OWNED_TABLES`` / ``env.include_object``. +__all__ = ["LANGGRAPH_OWNED_TABLES", "include_object"] # Import all models so metadata is populated. try: @@ -39,6 +54,7 @@ def run_migrations_offline() -> None: target_metadata=target_metadata, literal_binds=True, render_as_batch=True, + include_object=include_object, ) with context.begin_transaction(): context.run_migrations() @@ -49,6 +65,7 @@ def do_run_migrations(connection): connection=connection, target_metadata=target_metadata, render_as_batch=True, # Required for SQLite ALTER TABLE support + include_object=include_object, ) with context.begin_transaction(): context.run_migrations() @@ -56,6 +73,25 @@ def do_run_migrations(connection): async def run_migrations_online() -> None: connectable = create_async_engine(config.get_main_option("sqlalchemy.url")) + + # Cross-process bootstrap safety for SQLite: every connection alembic + # opens needs a wide ``busy_timeout`` so that when another process holds + # the file write lock (e.g. mid-bootstrap), our writes wait instead of + # raising ``database is locked``. The production engine in + # ``deerflow.persistence.engine`` sets this on its own connections, but + # alembic spawns its OWN engine here -- those connections wouldn't inherit + # anything unless we wire the same hook on this one. + if connectable.url.drivername.startswith("sqlite"): + from sqlalchemy import event + + @event.listens_for(connectable.sync_engine, "connect") + def _alembic_sqlite_busy_timeout(dbapi_conn, _record): # noqa: ARG001 + cursor = dbapi_conn.cursor() + try: + cursor.execute("PRAGMA busy_timeout=30000;") + finally: + cursor.close() + async with connectable.connect() as connection: await connection.run_sync(do_run_migrations) await connectable.dispose() diff --git a/backend/packages/harness/deerflow/persistence/migrations/script.py.mako b/backend/packages/harness/deerflow/persistence/migrations/script.py.mako new file mode 100644 index 000000000..a3254f9db --- /dev/null +++ b/backend/packages/harness/deerflow/persistence/migrations/script.py.mako @@ -0,0 +1,31 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} + +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision: str = ${repr(up_revision)} +down_revision: str | Sequence[str] | None = ${repr(down_revision)} +branch_labels: str | Sequence[str] | None = ${repr(branch_labels)} +depends_on: str | Sequence[str] | None = ${repr(depends_on)} + + +def upgrade() -> None: + """Upgrade schema.""" + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + """Downgrade schema.""" + ${downgrades if downgrades else "pass"} diff --git a/backend/packages/harness/deerflow/persistence/migrations/versions/0001_baseline.py b/backend/packages/harness/deerflow/persistence/migrations/versions/0001_baseline.py new file mode 100644 index 000000000..f50281885 --- /dev/null +++ b/backend/packages/harness/deerflow/persistence/migrations/versions/0001_baseline.py @@ -0,0 +1,291 @@ +"""baseline -- chain root for DeerFlow application schema. + +Revision ID: 0001_baseline +Revises: +Create Date: 2026-06-22 + +Role of this revision +===================== + +This revision encodes the schema that ``Base.metadata.create_all`` produces for +every DeerFlow-owned table at the point alembic was wired in. Under the hybrid +bootstrap strategy (``deerflow.persistence.bootstrap.bootstrap_schema``), the +``upgrade()`` here is **almost never executed**: + +- Fresh DB -> ``create_all`` + ``alembic stamp head`` (no upgrade run). +- Legacy DB -> ``alembic stamp 0001_baseline`` + ``upgrade head`` + (jumps directly to the next revision; baseline ``upgrade()`` + is also not run, because alembic only runs revisions strictly + AFTER the stamped position). +- Versioned DB -> ``upgrade head`` (continues from whatever revision is in + ``alembic_version``; baseline ``upgrade()`` only runs when + the DB happens to be at ``base``). + +The baseline therefore primarily serves as a **stamp target + chain root**. +``upgrade()`` is kept faithful to ``Base.metadata`` so ``alembic upgrade base +-> head`` round-trips in test fixtures and ``downgrade()`` is provided in full +for symmetry, but production-path correctness does not depend on this +revision's DDL matching ``create_all`` byte-for-byte. + +LangGraph checkpointer tables (``checkpoints``, ``checkpoint_blobs``, +``checkpoint_writes``, ``checkpoint_migrations``) are intentionally absent -- +they belong to LangGraph and are excluded by ``env.py::include_object``. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "0001_baseline" +down_revision: str | Sequence[str] | None = None +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + """Upgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "channel_connections", + sa.Column("id", sa.String(length=64), nullable=False), + sa.Column("owner_user_id", sa.String(length=64), nullable=False), + sa.Column("provider", sa.String(length=32), nullable=False), + sa.Column("status", sa.String(length=32), nullable=False), + sa.Column("external_account_id", sa.String(length=128), nullable=False), + sa.Column("external_account_name", sa.String(length=256), nullable=True), + sa.Column("workspace_id", sa.String(length=128), nullable=False), + sa.Column("workspace_name", sa.String(length=256), nullable=True), + sa.Column("bot_user_id", sa.String(length=128), nullable=True), + sa.Column("scopes_json", sa.JSON(), nullable=False), + sa.Column("capabilities_json", sa.JSON(), nullable=False), + sa.Column("metadata_json", sa.JSON(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("last_seen_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("last_error_at", sa.DateTime(timezone=True), nullable=True), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("owner_user_id", "provider", "external_account_id", "workspace_id", name="uq_channel_connection_owner_provider_identity"), + ) + with op.batch_alter_table("channel_connections", schema=None) as batch_op: + batch_op.create_index("idx_channel_connections_event_lookup", ["provider", "workspace_id", "bot_user_id"], unique=False) + batch_op.create_index(batch_op.f("ix_channel_connections_owner_user_id"), ["owner_user_id"], unique=False) + batch_op.create_index(batch_op.f("ix_channel_connections_provider"), ["provider"], unique=False) + batch_op.create_index("uq_channel_connection_active_identity", ["provider", "external_account_id", "workspace_id"], unique=True, sqlite_where=sa.text("status != 'revoked'"), postgresql_where=sa.text("status != 'revoked'")) + + op.create_table( + "channel_oauth_states", + sa.Column("state_hash", sa.String(length=128), nullable=False), + sa.Column("owner_user_id", sa.String(length=64), nullable=False), + sa.Column("provider", sa.String(length=32), nullable=False), + sa.Column("code_verifier_encrypted", sa.Text(), nullable=True), + sa.Column("nonce_hash", sa.String(length=128), nullable=True), + sa.Column("redirect_after", sa.Text(), nullable=True), + sa.Column("requested_scopes_json", sa.JSON(), nullable=False), + sa.Column("metadata_json", sa.JSON(), nullable=False), + sa.Column("expires_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("consumed_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("state_hash"), + ) + with op.batch_alter_table("channel_oauth_states", schema=None) as batch_op: + batch_op.create_index(batch_op.f("ix_channel_oauth_states_owner_user_id"), ["owner_user_id"], unique=False) + batch_op.create_index(batch_op.f("ix_channel_oauth_states_provider"), ["provider"], unique=False) + + op.create_table( + "feedback", + sa.Column("feedback_id", sa.String(length=64), nullable=False), + sa.Column("run_id", sa.String(length=64), nullable=False), + sa.Column("thread_id", sa.String(length=64), nullable=False), + sa.Column("user_id", sa.String(length=64), nullable=True), + sa.Column("message_id", sa.String(length=64), nullable=True), + sa.Column("rating", sa.Integer(), nullable=False), + sa.Column("comment", sa.Text(), nullable=True), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("feedback_id"), + sa.UniqueConstraint("thread_id", "run_id", "user_id", name="uq_feedback_thread_run_user"), + ) + with op.batch_alter_table("feedback", schema=None) as batch_op: + batch_op.create_index(batch_op.f("ix_feedback_run_id"), ["run_id"], unique=False) + batch_op.create_index(batch_op.f("ix_feedback_thread_id"), ["thread_id"], unique=False) + batch_op.create_index(batch_op.f("ix_feedback_user_id"), ["user_id"], unique=False) + + op.create_table( + "run_events", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("thread_id", sa.String(length=64), nullable=False), + sa.Column("run_id", sa.String(length=64), nullable=False), + sa.Column("user_id", sa.String(length=64), nullable=True), + sa.Column("event_type", sa.String(length=32), nullable=False), + sa.Column("category", sa.String(length=16), nullable=False), + sa.Column("content", sa.Text(), nullable=False), + sa.Column("event_metadata", sa.JSON(), nullable=False), + sa.Column("seq", sa.Integer(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("thread_id", "seq", name="uq_events_thread_seq"), + ) + with op.batch_alter_table("run_events", schema=None) as batch_op: + batch_op.create_index("ix_events_run", ["thread_id", "run_id", "seq"], unique=False) + batch_op.create_index("ix_events_thread_cat_seq", ["thread_id", "category", "seq"], unique=False) + batch_op.create_index(batch_op.f("ix_run_events_user_id"), ["user_id"], unique=False) + + op.create_table( + "runs", + sa.Column("run_id", sa.String(length=64), nullable=False), + sa.Column("thread_id", sa.String(length=64), nullable=False), + sa.Column("assistant_id", sa.String(length=128), nullable=True), + sa.Column("user_id", sa.String(length=64), nullable=True), + sa.Column("status", sa.String(length=20), nullable=False), + sa.Column("model_name", sa.String(length=128), nullable=True), + sa.Column("multitask_strategy", sa.String(length=20), nullable=False), + sa.Column("metadata_json", sa.JSON(), nullable=False), + sa.Column("kwargs_json", sa.JSON(), nullable=False), + sa.Column("error", sa.Text(), nullable=True), + sa.Column("message_count", sa.Integer(), nullable=False), + sa.Column("first_human_message", sa.Text(), nullable=True), + sa.Column("last_ai_message", sa.Text(), nullable=True), + sa.Column("total_input_tokens", sa.Integer(), nullable=False), + sa.Column("total_output_tokens", sa.Integer(), nullable=False), + sa.Column("total_tokens", sa.Integer(), nullable=False), + sa.Column("llm_call_count", sa.Integer(), nullable=False), + sa.Column("lead_agent_tokens", sa.Integer(), nullable=False), + sa.Column("subagent_tokens", sa.Integer(), nullable=False), + sa.Column("middleware_tokens", sa.Integer(), nullable=False), + sa.Column("token_usage_by_model", sa.JSON(), nullable=False, server_default=sa.text("'{}'")), + sa.Column("follow_up_to_run_id", sa.String(length=64), nullable=True), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("run_id"), + ) + with op.batch_alter_table("runs", schema=None) as batch_op: + batch_op.create_index(batch_op.f("ix_runs_thread_id"), ["thread_id"], unique=False) + batch_op.create_index("ix_runs_thread_status", ["thread_id", "status"], unique=False) + batch_op.create_index(batch_op.f("ix_runs_user_id"), ["user_id"], unique=False) + + op.create_table( + "threads_meta", + sa.Column("thread_id", sa.String(length=64), nullable=False), + sa.Column("assistant_id", sa.String(length=128), nullable=True), + sa.Column("user_id", sa.String(length=64), nullable=True), + sa.Column("display_name", sa.String(length=256), nullable=True), + sa.Column("status", sa.String(length=20), nullable=False), + sa.Column("metadata_json", sa.JSON(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("thread_id"), + ) + with op.batch_alter_table("threads_meta", schema=None) as batch_op: + batch_op.create_index(batch_op.f("ix_threads_meta_assistant_id"), ["assistant_id"], unique=False) + batch_op.create_index(batch_op.f("ix_threads_meta_user_id"), ["user_id"], unique=False) + + op.create_table( + "users", + sa.Column("id", sa.String(length=36), nullable=False), + sa.Column("email", sa.String(length=320), nullable=False), + sa.Column("password_hash", sa.String(length=128), nullable=True), + sa.Column("system_role", sa.String(length=16), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("oauth_provider", sa.String(length=32), nullable=True), + sa.Column("oauth_id", sa.String(length=128), nullable=True), + sa.Column("needs_setup", sa.Boolean(), nullable=False), + sa.Column("token_version", sa.Integer(), nullable=False), + sa.PrimaryKeyConstraint("id"), + ) + with op.batch_alter_table("users", schema=None) as batch_op: + batch_op.create_index("idx_users_oauth_identity", ["oauth_provider", "oauth_id"], unique=True, sqlite_where=sa.text("oauth_provider IS NOT NULL AND oauth_id IS NOT NULL")) + batch_op.create_index(batch_op.f("ix_users_email"), ["email"], unique=True) + + op.create_table( + "channel_conversations", + sa.Column("id", sa.String(length=64), nullable=False), + sa.Column("connection_id", sa.String(length=64), nullable=False), + sa.Column("owner_user_id", sa.String(length=64), nullable=False), + sa.Column("provider", sa.String(length=32), nullable=False), + sa.Column("external_conversation_id", sa.String(length=128), nullable=False), + sa.Column("external_topic_id", sa.String(length=128), nullable=False), + sa.Column("thread_id", sa.String(length=64), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.ForeignKeyConstraint(["connection_id"], ["channel_connections.id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("connection_id", "external_conversation_id", "external_topic_id", name="uq_channel_conversation_connection_external"), + ) + with op.batch_alter_table("channel_conversations", schema=None) as batch_op: + batch_op.create_index(batch_op.f("ix_channel_conversations_connection_id"), ["connection_id"], unique=False) + batch_op.create_index(batch_op.f("ix_channel_conversations_owner_user_id"), ["owner_user_id"], unique=False) + batch_op.create_index(batch_op.f("ix_channel_conversations_provider"), ["provider"], unique=False) + batch_op.create_index(batch_op.f("ix_channel_conversations_thread_id"), ["thread_id"], unique=False) + + op.create_table( + "channel_credentials", + sa.Column("connection_id", sa.String(length=64), nullable=False), + sa.Column("encrypted_access_token", sa.Text(), nullable=True), + sa.Column("encrypted_refresh_token", sa.Text(), nullable=True), + sa.Column("token_type", sa.String(length=32), nullable=True), + sa.Column("expires_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("refresh_expires_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("encrypted_extra_json", sa.Text(), nullable=True), + sa.Column("version", sa.Integer(), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.ForeignKeyConstraint(["connection_id"], ["channel_connections.id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("connection_id"), + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + """Downgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table("channel_credentials") + with op.batch_alter_table("channel_conversations", schema=None) as batch_op: + batch_op.drop_index(batch_op.f("ix_channel_conversations_thread_id")) + batch_op.drop_index(batch_op.f("ix_channel_conversations_provider")) + batch_op.drop_index(batch_op.f("ix_channel_conversations_owner_user_id")) + batch_op.drop_index(batch_op.f("ix_channel_conversations_connection_id")) + + op.drop_table("channel_conversations") + with op.batch_alter_table("users", schema=None) as batch_op: + batch_op.drop_index(batch_op.f("ix_users_email")) + batch_op.drop_index("idx_users_oauth_identity", sqlite_where=sa.text("oauth_provider IS NOT NULL AND oauth_id IS NOT NULL")) + + op.drop_table("users") + with op.batch_alter_table("threads_meta", schema=None) as batch_op: + batch_op.drop_index(batch_op.f("ix_threads_meta_user_id")) + batch_op.drop_index(batch_op.f("ix_threads_meta_assistant_id")) + + op.drop_table("threads_meta") + with op.batch_alter_table("runs", schema=None) as batch_op: + batch_op.drop_index(batch_op.f("ix_runs_user_id")) + batch_op.drop_index("ix_runs_thread_status") + batch_op.drop_index(batch_op.f("ix_runs_thread_id")) + + op.drop_table("runs") + with op.batch_alter_table("run_events", schema=None) as batch_op: + batch_op.drop_index(batch_op.f("ix_run_events_user_id")) + batch_op.drop_index("ix_events_thread_cat_seq") + batch_op.drop_index("ix_events_run") + + op.drop_table("run_events") + with op.batch_alter_table("feedback", schema=None) as batch_op: + batch_op.drop_index(batch_op.f("ix_feedback_user_id")) + batch_op.drop_index(batch_op.f("ix_feedback_thread_id")) + batch_op.drop_index(batch_op.f("ix_feedback_run_id")) + + op.drop_table("feedback") + with op.batch_alter_table("channel_oauth_states", schema=None) as batch_op: + batch_op.drop_index(batch_op.f("ix_channel_oauth_states_provider")) + batch_op.drop_index(batch_op.f("ix_channel_oauth_states_owner_user_id")) + + op.drop_table("channel_oauth_states") + with op.batch_alter_table("channel_connections", schema=None) as batch_op: + batch_op.drop_index("uq_channel_connection_active_identity", sqlite_where=sa.text("status != 'revoked'"), postgresql_where=sa.text("status != 'revoked'")) + batch_op.drop_index(batch_op.f("ix_channel_connections_provider")) + batch_op.drop_index(batch_op.f("ix_channel_connections_owner_user_id")) + batch_op.drop_index("idx_channel_connections_event_lookup") + + op.drop_table("channel_connections") + # ### end Alembic commands ### diff --git a/backend/packages/harness/deerflow/persistence/migrations/versions/0002_runs_token_usage.py b/backend/packages/harness/deerflow/persistence/migrations/versions/0002_runs_token_usage.py new file mode 100644 index 000000000..f4bf6fa1c --- /dev/null +++ b/backend/packages/harness/deerflow/persistence/migrations/versions/0002_runs_token_usage.py @@ -0,0 +1,69 @@ +"""Add ``runs.token_usage_by_model`` column. + +Revision ID: 0002_runs_token_usage +Revises: 0001_baseline +Create Date: 2026-06-22 + +Fixes GitHub issue #3682: any pre-existing DB (created before commit e7a03e52 +on PR #3658) lacks the ``token_usage_by_model`` JSON column on ``runs``. +Without this migration, every endpoint that ``SELECT``s from ``runs`` raises +``no such column: runs.token_usage_by_model``. + +Schema parity with ``Base.metadata`` +------------------------------------ + +The ORM model declares the column as ``Mapped[dict] = mapped_column(JSON, +default=dict, server_default=text("'{}'"))`` -- non-Optional, so SQLAlchemy +infers ``nullable=False``. ``Base.metadata.create_all`` (the empty-DB +bootstrap path) therefore produces ``token_usage_by_model JSON NOT NULL +DEFAULT '{}'`` on fresh databases. + +To keep legacy-upgraded databases schema-identical to fresh ones, this +migration adds the column with the same ``nullable=False`` and +``server_default='{}'``. The server default is also what lets +``ALTER TABLE runs ADD COLUMN ... NOT NULL`` succeed on a populated table: +existing rows pick up the empty-object default at ALTER time instead of +triggering ``NOT NULL`` violations. + +Idempotency +----------- + +Uses ``safe_add_column`` so re-running this revision against a DB where the +column already exists is a no-op. That covers two real cases: + +1. Users who applied the workaround in the issue manually + (``ALTER TABLE runs ADD COLUMN token_usage_by_model JSON``). +2. Concurrent bootstrap on multiple Gateway instances if the cross-process + lock is somehow bypassed -- defence-in-depth on top of + ``bootstrap_schema``'s advisory-lock / sentinel-row mutex. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import sqlalchemy as sa + +from deerflow.persistence.migrations._helpers import safe_add_column, safe_drop_column + +# revision identifiers, used by Alembic. +revision: str = "0002_runs_token_usage" +down_revision: str | Sequence[str] | None = "0001_baseline" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + safe_add_column( + "runs", + sa.Column( + "token_usage_by_model", + sa.JSON(), + nullable=False, + server_default=sa.text("'{}'"), + ), + ) + + +def downgrade() -> None: + safe_drop_column("runs", "token_usage_by_model") diff --git a/backend/packages/harness/deerflow/persistence/run/model.py b/backend/packages/harness/deerflow/persistence/run/model.py index 4fa7fc3ad..1d5f16f48 100644 --- a/backend/packages/harness/deerflow/persistence/run/model.py +++ b/backend/packages/harness/deerflow/persistence/run/model.py @@ -4,7 +4,7 @@ from __future__ import annotations from datetime import UTC, datetime -from sqlalchemy import JSON, DateTime, Index, String, Text +from sqlalchemy import JSON, DateTime, Index, String, Text, text from sqlalchemy.orm import Mapped, mapped_column from deerflow.persistence.base import Base @@ -39,8 +39,7 @@ class RunRow(Base): lead_agent_tokens: Mapped[int] = mapped_column(default=0) subagent_tokens: Mapped[int] = mapped_column(default=0) middleware_tokens: Mapped[int] = mapped_column(default=0) - # Per-model token breakdown - token_usage_by_model: Mapped[dict] = mapped_column(JSON, default=dict) + token_usage_by_model: Mapped[dict] = mapped_column(JSON, default=dict, server_default=text("'{}'")) # Follow-up association follow_up_to_run_id: Mapped[str | None] = mapped_column(String(64)) diff --git a/backend/scripts/_autogen_revision.py b/backend/scripts/_autogen_revision.py new file mode 100644 index 000000000..90b098088 --- /dev/null +++ b/backend/scripts/_autogen_revision.py @@ -0,0 +1,78 @@ +"""Generate a new alembic revision against an ephemeral SQLite DB. + +Used by ``make migrate-rev MSG="..."``. Avoids two pitfalls: + +1. ``alembic.ini``'s default ``sqlalchemy.url`` (``sqlite:///./data/deerflow.db``) + points at a path that doesn't exist in a clean checkout, so a bare + ``alembic revision --autogenerate`` fails with ``unable to open database file``. +2. A persistent DB might be at an unknown revision (or at no revision at all), + producing a noisy autogenerate diff that mixes "real" changes with + accidentally-detected drift. + +This script builds a *fresh* temp SQLite, runs the existing alembic chain to +``head`` against it, then runs ``alembic revision --autogenerate`` against +that. The temp DB must be built from migration history -- not from +``Base.metadata.create_all`` -- so newly edited ORM fields that do not yet have +a revision remain visible to autogenerate as a real diff. + +The generated file lands in +``packages/harness/deerflow/persistence/migrations/versions/`` -- exactly +where alembic puts it by default -- and the temp directory is left for the OS +to GC. Review the generated revision and switch raw ``op.add_column`` / +``op.drop_column`` calls to the idempotent helpers in ``migrations/_helpers.py`` +before committing. + +Run from the ``backend/`` directory: + PYTHONPATH=. uv run python scripts/_autogen_revision.py "MESSAGE" +or via Makefile: + make migrate-rev MSG="..." +""" + +from __future__ import annotations + +import os +import sys +import tempfile +from pathlib import Path + +from alembic import command +from alembic.config import Config + +import deerflow.persistence.models # noqa: F401 -- registers ORM models with Base.metadata +from deerflow.persistence.bootstrap import _escape_url_for_alembic + +BACKEND_DIR = Path(__file__).resolve().parents[1] +MIGRATIONS_DIR = BACKEND_DIR / "packages/harness/deerflow/persistence/migrations" + + +def _alembic_config(url: str) -> Config: + cfg = Config() + cfg.set_main_option("script_location", str(MIGRATIONS_DIR)) + # Shared with ``bootstrap._alembic_safe_url`` so the ConfigParser ``%`` + # interpolation rule lives in one place. + cfg.set_main_option("sqlalchemy.url", _escape_url_for_alembic(url)) + return cfg + + +def _build_temp_db_at_head() -> str: + tmpdir = tempfile.mkdtemp(prefix="deerflow-autogen-") + db_path = os.path.join(tmpdir, "autogen.db").replace(os.sep, "/") + url = f"sqlite+aiosqlite:///{db_path}" + command.upgrade(_alembic_config(url), "head") + return url + + +def main() -> None: + if len(sys.argv) < 2 or not sys.argv[1].strip(): + print('usage: python scripts/_autogen_revision.py "describe the change"', file=sys.stderr) + sys.exit(2) + message = sys.argv[1] + + url = _build_temp_db_at_head() + print(f"autogen: built temp DB at head: {url}", file=sys.stderr) + + command.revision(_alembic_config(url), message=message, autogenerate=True) + + +if __name__ == "__main__": + main() diff --git a/backend/tests/blocking_io/test_persistence_bootstrap.py b/backend/tests/blocking_io/test_persistence_bootstrap.py new file mode 100644 index 000000000..ff81d1be2 --- /dev/null +++ b/backend/tests/blocking_io/test_persistence_bootstrap.py @@ -0,0 +1,75 @@ +"""Regression: ``bootstrap_schema`` offloads ``alembic.command.stamp`` / +``alembic.command.upgrade`` via ``asyncio.to_thread``. + +The alembic commands are synchronous: they open their own engine and execute +DDL. Calling them directly on the FastAPI lifespan event loop would block -- +exactly the failure mode of the issue chain that motivated the hybrid +bootstrap (sync IO on the loop = silent stalls / timeouts). + +Anchor strategy +--------------- + +We can't run a real ``init_engine(backend="sqlite", ...)`` under the strict +Blockbuster gate without tripping on ``create_async_engine``'s own +``os.path.abspath`` (which is a pre-existing concern, not the bootstrap's). +The companion ``test_persistence_engine_sqlite.py`` covers the ``init_engine`` +makedirs offload by mocking ``create_async_engine`` away entirely. That same +mocking approach would defeat the point here, because the alembic stamp / +upgrade calls in ``bootstrap_schema`` need a *real* on-disk SQLite DB to +exercise. + +So this test installs a spy on ``asyncio.to_thread`` and confirms that the +two alembic entry points -- ``_stamp`` and ``_upgrade`` from +``bootstrap_schema`` -- are dispatched through it, not invoked inline. If a +future refactor inlines either call, the spy records zero invocations for +that function and the assertion fails. +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path + +import pytest +from sqlalchemy.ext.asyncio import create_async_engine + +import deerflow.persistence.models # noqa: F401 +from deerflow.persistence import bootstrap as bootstrap_mod + +pytestmark = pytest.mark.asyncio + + +@pytest.mark.allow_blocking_io +async def test_bootstrap_offloads_alembic_stamp_and_upgrade(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """Stamp + upgrade must go through ``asyncio.to_thread``. + + Marked ``allow_blocking_io`` so the strict Blockbuster gate does not flag + incidental blocking IO in test-fixture setup (engine creation paths, + SQLite path resolution). The point of this test is the + ``asyncio.to_thread`` wrapping invariant, which the spy below checks + deterministically. + """ + seen: list[str] = [] + + original_to_thread = asyncio.to_thread + + async def spy_to_thread(func, *args, **kwargs): + seen.append(getattr(func, "__name__", repr(func))) + return await original_to_thread(func, *args, **kwargs) + + monkeypatch.setattr(bootstrap_mod.asyncio, "to_thread", spy_to_thread) + + # Use a real SQLite DB so alembic actually runs stamp + upgrade. + db_path = tmp_path / "spy.db" + engine = create_async_engine(f"sqlite+aiosqlite:///{db_path.as_posix()}") + try: + # Empty branch -> create_all + stamp head. ``_stamp`` must be offloaded. + await bootstrap_mod.bootstrap_schema(engine, backend="sqlite") + assert "_stamp" in seen, f"_stamp not offloaded; saw: {seen}" + + # Re-run -> versioned branch -> upgrade head (no-op at head). ``_upgrade`` must be offloaded. + seen.clear() + await bootstrap_mod.bootstrap_schema(engine, backend="sqlite") + assert "_upgrade" in seen, f"_upgrade not offloaded; saw: {seen}" + finally: + await engine.dispose() diff --git a/backend/tests/blocking_io/test_persistence_engine_sqlite.py b/backend/tests/blocking_io/test_persistence_engine_sqlite.py index f7cc3daa8..3b1ff1747 100644 --- a/backend/tests/blocking_io/test_persistence_engine_sqlite.py +++ b/backend/tests/blocking_io/test_persistence_engine_sqlite.py @@ -11,6 +11,12 @@ Blockbuster context with a `sqlite_dir` that does not yet exist, so `os.makedirs actually runs. The async engine/session machinery is mocked out so the only host filesystem operation under test is the directory creation; if it regresses to run directly on the event loop, Blockbuster raises `BlockingError` and this fails. + +We also stub ``bootstrap_schema`` so the alembic stamp/upgrade path -- which has +its own ``asyncio.to_thread`` regression anchor in +``test_persistence_bootstrap.py`` -- does not turn this test into a +double-coverage one. Keeping concerns separated means a regression in either +offload (makedirs vs alembic) points at the right place. """ from __future__ import annotations @@ -49,10 +55,17 @@ async def test_init_engine_sqlite_dir_setup_does_not_block_event_loop(tmp_path: mock_engine.begin.return_value = begin_ctx mock_engine.dispose = AsyncMock() + async def _noop_bootstrap(*_args, **_kwargs): + return None + with ( patch.object(engine_mod, "create_async_engine", return_value=mock_engine), patch.object(engine_mod, "async_sessionmaker", return_value=MagicMock()), patch("sqlalchemy.event.listens_for", _noop_listens_for), + patch( + "deerflow.persistence.bootstrap.bootstrap_schema", + new=_noop_bootstrap, + ), ): await engine_mod.init_engine( backend="sqlite", diff --git a/backend/tests/test_base_to_dict.py b/backend/tests/test_base_to_dict.py index a01ca8b5a..1e3dbe841 100644 --- a/backend/tests/test_base_to_dict.py +++ b/backend/tests/test_base_to_dict.py @@ -7,7 +7,7 @@ must stay identical. from __future__ import annotations -from sqlalchemy import Integer, String +from sqlalchemy import Integer, MetaData, String from sqlalchemy.orm import Mapped, mapped_column from deerflow.persistence.base import Base, _column_keys @@ -15,6 +15,10 @@ from deerflow.persistence.base import Base, _column_keys class _Widget(Base): __tablename__ = "_widget_to_dict_test" + # Keep this test-only model out of the application metadata. Pytest imports + # test modules during collection, so registering it on ``Base.metadata`` + # would leak the table into unrelated create_all/schema-parity tests. + metadata = MetaData() id: Mapped[int] = mapped_column(Integer, primary_key=True) name: Mapped[str] = mapped_column(String(32)) diff --git a/backend/tests/test_persistence_autogen_script.py b/backend/tests/test_persistence_autogen_script.py new file mode 100644 index 000000000..d35fdb926 --- /dev/null +++ b/backend/tests/test_persistence_autogen_script.py @@ -0,0 +1,102 @@ +"""Tests for ``scripts/_autogen_revision.py`` (``make migrate-rev``). + +The script must work in a clean checkout without any pre-existing data +directory -- this is the failure mode reported as P2: a bare ``alembic +revision --autogenerate`` would crash with +``sqlite3.OperationalError: unable to open database file`` because +``alembic.ini``'s default URL points at ``./data/deerflow.db`` which doesn't +exist yet. + +The fix: the script builds its own temp DB by running the existing alembic +chain to head and runs autogenerate against THAT, instead of relying on +``alembic.ini``'s URL or runtime ``create_all`` bootstrap. +""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest +import sqlalchemy as sa + +import deerflow.persistence.models # noqa: F401 +from deerflow.persistence.base import Base + + +@pytest.fixture(scope="module") +def autogen_module(): + """Load ``scripts/_autogen_revision.py`` as an importable module. + + The file lives outside the package tree (under ``backend/scripts/``) so we + load it directly via ``spec_from_file_location``. + """ + script_path = Path(__file__).resolve().parents[1] / "scripts/_autogen_revision.py" + assert script_path.exists(), f"missing autogen script at {script_path}" + spec = importlib.util.spec_from_file_location("_autogen_revision_under_test", script_path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_autogen_builds_temp_db_at_head_without_data_dir(autogen_module, monkeypatch) -> None: + """The temp-DB builder must succeed even when ``./data/`` does not exist. + + We chdir to an empty directory to mimic a clean checkout where the + alembic.ini default URL would explode. + """ + import os # noqa: PLC0415 + import tempfile # noqa: PLC0415 + + workdir = tempfile.mkdtemp(prefix="deerflow-autogen-test-") + monkeypatch.chdir(workdir) + # Sanity: this directory has no ``./data/`` -- so the alembic.ini default + # URL would fail if used. + assert not os.path.exists("data") + + url = autogen_module._build_temp_db_at_head() + assert url.startswith("sqlite+aiosqlite:///"), f"unexpected URL shape: {url}" + # The temp DB file should now exist. + db_path = url.replace("sqlite+aiosqlite:///", "") + assert os.path.exists(db_path), f"temp DB file not created at {db_path}" + + +def test_autogen_temp_db_is_at_head(autogen_module) -> None: + """The temp DB the autogen script builds must be at head, so the + autogenerate diff against current models is empty (or only reflects + intentional, in-progress model changes).""" + import sqlite3 # noqa: PLC0415 + + url = autogen_module._build_temp_db_at_head() + db_path = url.replace("sqlite+aiosqlite:///", "") + with sqlite3.connect(db_path) as raw: + row = raw.execute("SELECT version_num FROM alembic_version").fetchone() + assert row is not None, "autogen temp DB has no alembic_version row -- bootstrap failed" + # head is whatever the script tree currently says; we just assert it's there. + assert row[0] + + +def test_autogen_temp_db_comes_from_migration_history_not_current_metadata(autogen_module) -> None: + """Pending ORM changes must remain visible to autogenerate. + + If the helper accidentally uses runtime ``bootstrap_schema`` / + ``Base.metadata.create_all`` again, this probe table would be created in + the temp DB and the test would fail. A temp DB built from alembic history + only contains objects that committed revisions know how to create. + """ + import sqlite3 # noqa: PLC0415 + + probe_name = "__autogen_probe_pending_migration__" + probe_table = sa.Table(probe_name, Base.metadata, sa.Column("id", sa.Integer, primary_key=True)) + try: + url = autogen_module._build_temp_db_at_head() + db_path = url.replace("sqlite+aiosqlite:///", "") + with sqlite3.connect(db_path) as raw: + exists = raw.execute( + "SELECT count(*) FROM sqlite_master WHERE type='table' AND name=?", + (probe_name,), + ).fetchone()[0] + assert exists == 0, "temp DB was built from current ORM metadata instead of migration history" + finally: + Base.metadata.remove(probe_table) diff --git a/backend/tests/test_persistence_bootstrap.py b/backend/tests/test_persistence_bootstrap.py new file mode 100644 index 000000000..8d3a1a67e --- /dev/null +++ b/backend/tests/test_persistence_bootstrap.py @@ -0,0 +1,636 @@ +"""Tests for ``deerflow.persistence.bootstrap.bootstrap_schema``. + +Covers the three-branch decision table: + +| DB state | Action | +|---------------------------------------|-----------------------------------------| +| empty | create_all + stamp head | +| legacy (DeerFlow tables, no alembic_version) | create_all (baseline tables only, backfill) + stamp baseline + upgrade head | +| versioned | upgrade head | + +Each test seeds a temp SQLite to the relevant pre-state, runs +``bootstrap_schema``, and asserts both the resulting schema and the +``alembic_version`` row. + +The legacy branch is exercised across three scenarios: token-usage column +missing, token-usage column already present, and a baseline-era table +missing entirely (the ``channel_*`` backfill case). The first two prove the +column-level idempotent helpers handle both sub-cases; the third proves the +table-level backfill works. +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path + +import pytest +import sqlalchemy as sa +from sqlalchemy.ext.asyncio import create_async_engine + +# Pre-import models so Base.metadata is populated before bootstrap reads it. +import deerflow.persistence.models # noqa: F401 +from deerflow.persistence.base import Base +from deerflow.persistence.bootstrap import ( + _BASELINE_TABLE_NAMES, + _decide_state, + _get_alembic_config, + _get_head_revision, + _run_baseline_create_all_sync, + _upgrade, + bootstrap_schema, +) +from deerflow.persistence.migrations._helpers import _normalize_default + +# Mark only async tests via the decorator below; module-level pytestmark would +# spuriously warn for the sync ``TestDecideState`` cases. +asyncio_test = pytest.mark.asyncio + + +HEAD = "0002_runs_token_usage" +BASELINE = "0001_baseline" + + +def _url(tmp_path: Path, name: str = "test.db") -> str: + return f"sqlite+aiosqlite:///{(tmp_path / name).as_posix()}" + + +async def _table_names(engine) -> set[str]: + async with engine.connect() as conn: + return await conn.run_sync(lambda c: set(sa.inspect(c).get_table_names())) + + +async def _runs_columns(engine) -> set[str]: + async with engine.connect() as conn: + return await conn.run_sync(lambda c: {col["name"] for col in sa.inspect(c).get_columns("runs")}) + + +async def _runs_column_meta(engine, column_name: str) -> dict: + async with engine.connect() as conn: + cols = await conn.run_sync(lambda c: sa.inspect(c).get_columns("runs")) + for c in cols: + if c["name"] == column_name: + return c + raise AssertionError(f"column {column_name!r} not found in runs") + + +async def _alembic_version(engine) -> str | None: + async with engine.connect() as conn: + row = await conn.execute(sa.text("SELECT version_num FROM alembic_version")) + return row.scalar() + + +async def _seed_legacy_without_column(engine) -> None: + """Build the pre-#3658 schema: create_all, then drop the new column.""" + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + async with engine.begin() as conn: + # SQLite supports DROP COLUMN from 3.35.0; the test runner pins recent + # Python which bundles a 3.40+ sqlite, so this is safe. + await conn.execute(sa.text("ALTER TABLE runs DROP COLUMN token_usage_by_model")) + + +async def _seed_legacy_with_column(engine) -> None: + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + +async def _seed_legacy_missing_channel_tables(engine) -> None: + """Build a pre-#1930 schema: baseline tables exist but ``channel_*`` do not. + + Models the worst-case legacy DB the bootstrap layer has to repair -- a + user who upgraded across multiple releases and never had the channel_* + tables provisioned in the first place. We achieve it by running the full + ``create_all`` and then dropping the channel_* tables in FK-dependency + order (credentials/conversations reference channel_connections). + """ + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + async with engine.begin() as conn: + for table in ( + "channel_credentials", + "channel_conversations", + "channel_oauth_states", + "channel_connections", + ): + await conn.execute(sa.text(f"DROP TABLE IF EXISTS {table}")) + + +# --------------------------------------------------------------------------- +# Branch 1: empty DB +# --------------------------------------------------------------------------- + + +@asyncio_test +async def test_empty_branch_creates_all_and_stamps_head(tmp_path: Path) -> None: + engine = create_async_engine(_url(tmp_path)) + try: + await bootstrap_schema(engine, backend="sqlite") + tables = await _table_names(engine) + for required in { + "runs", + "threads_meta", + "feedback", + "users", + "run_events", + "channel_connections", + "channel_credentials", + "channel_conversations", + "channel_oauth_states", + "alembic_version", + }: + assert required in tables, f"missing table: {required}" + assert "token_usage_by_model" in await _runs_columns(engine) + assert await _alembic_version(engine) == HEAD + finally: + await engine.dispose() + + +# --------------------------------------------------------------------------- +# Branch 2: legacy DB without token_usage_by_model +# --------------------------------------------------------------------------- + + +@asyncio_test +async def test_legacy_without_column_branch_upgrades(tmp_path: Path) -> None: + engine = create_async_engine(_url(tmp_path)) + try: + await _seed_legacy_without_column(engine) + assert "token_usage_by_model" not in await _runs_columns(engine) + assert "alembic_version" not in await _table_names(engine) + + await bootstrap_schema(engine, backend="sqlite") + + assert "token_usage_by_model" in await _runs_columns(engine) + assert await _alembic_version(engine) == HEAD + finally: + await engine.dispose() + + +# --------------------------------------------------------------------------- +# Legacy backfill: a DB that pre-dates a later-added baseline table (e.g. the +# ``channel_*`` tables from PR #1930) must end up with all baseline tables +# after bootstrap, otherwise the channels API 500s with ``no such table``. +# The fix runs ``create_all`` (idempotent) before ``stamp 0001_baseline`` so +# missing baseline tables are backfilled with their current ORM schema. +# --------------------------------------------------------------------------- + + +@asyncio_test +async def test_legacy_missing_channel_tables_get_backfilled(tmp_path: Path) -> None: + engine = create_async_engine(_url(tmp_path)) + try: + await _seed_legacy_missing_channel_tables(engine) + tables = await _table_names(engine) + # Sanity-check the seeded pre-state: ``runs`` triggers the legacy + # branch (has_deerflow_tables=True, no alembic_version) while the + # channel_* tables are absent. + assert "runs" in tables + assert "alembic_version" not in tables + for missing in { + "channel_connections", + "channel_credentials", + "channel_conversations", + "channel_oauth_states", + }: + assert missing not in tables, f"seed should not have {missing}" + + await bootstrap_schema(engine, backend="sqlite") + + tables = await _table_names(engine) + for required in { + "channel_connections", + "channel_credentials", + "channel_conversations", + "channel_oauth_states", + }: + assert required in tables, f"legacy backfill missed: {required}" + assert await _alembic_version(engine) == HEAD + finally: + await engine.dispose() + + +# --------------------------------------------------------------------------- +# Branch 3: legacy DB that ALREADY has the column (post-#3658 create_all, +# or user-applied manual ALTER). The branch is the same as the +# legacy-without-column case -- bootstrap stamps baseline and tries to +# upgrade. The idempotent revision helper (``safe_add_column``) silently +# skips when the column is present, so the schema does not change. +# --------------------------------------------------------------------------- + + +@asyncio_test +async def test_legacy_with_column_branch_upgrade_is_idempotent(tmp_path: Path) -> None: + engine = create_async_engine(_url(tmp_path)) + try: + await _seed_legacy_with_column(engine) + assert "token_usage_by_model" in await _runs_columns(engine) + assert "alembic_version" not in await _table_names(engine) + cols_before = await _runs_columns(engine) + + await bootstrap_schema(engine, backend="sqlite") + + cols_after = await _runs_columns(engine) + assert cols_after == cols_before, "idempotent upgrade should not alter schema" + assert await _alembic_version(engine) == HEAD + finally: + await engine.dispose() + + +# --------------------------------------------------------------------------- +# Drift-warning guard: a column re-added by a manual ALTER (e.g. the #3682 +# workaround) survives the legacy branch because ``safe_add_column`` is +# name-keyed, but the helper must ``logger.warning`` so the operator notices +# the residual nullability / server_default / type drift from the model. +# Two scenarios are pinned: (1) nullable JSON workaround -- nullability + +# server_default drift fire, type matches; (2) ``TEXT NOT NULL DEFAULT +# '{}'`` workaround -- only type drifts, must STILL fire thanks to the +# JSON/TEXT family check in ``_type_equivalent``. +# --------------------------------------------------------------------------- + + +@asyncio_test +async def test_legacy_with_manual_workaround_column_warns_on_drift( + tmp_path: Path, + caplog: pytest.LogCaptureFixture, +) -> None: + engine = create_async_engine(_url(tmp_path)) + try: + # Pre-#3658 schema with a workaround-style re-add: nullable JSON, + # no server default -- diverges from the model's NOT NULL DEFAULT '{}'. + # Type matches (JSON vs JSON) so the type-equivalence check stays quiet + # and the warning fires purely on nullability + server_default. + await _seed_legacy_without_column(engine) + async with engine.begin() as conn: + await conn.execute(sa.text("ALTER TABLE runs ADD COLUMN token_usage_by_model JSON")) + + with caplog.at_level("WARNING", logger="deerflow.persistence.migrations._helpers"): + await bootstrap_schema(engine, backend="sqlite") + + # Bootstrap still completes -- the helper does not block on drift. + assert await _alembic_version(engine) == HEAD + # And the manually-added column survives untouched (no auto-repair). + col = await _runs_column_meta(engine, "token_usage_by_model") + assert col["nullable"] is True + + drift_warnings = [r for r in caplog.records if r.levelname == "WARNING" and r.name == "deerflow.persistence.migrations._helpers" and "safe_add_column" in r.getMessage() and "token_usage_by_model" in r.getMessage()] + assert drift_warnings, "expected safe_add_column to warn about the drifted column" + msg = drift_warnings[0].getMessage() + assert "nullable" in msg + assert "server_default" in msg + # Type info is always echoed in the payload for triage context. + assert "actual_type=" in msg and "desired_type=" in msg, f"warning missing type info: {msg!r}" + # JSON ≈ JSON, so the equivalence check must NOT produce a "type" diff + # entry here -- that would be a false positive on the matching-type case. + assert "type actual=" not in msg, f"unexpected type drift on matching JSON column: {msg!r}" + finally: + await engine.dispose() + + +@asyncio_test +async def test_legacy_with_wrong_type_workaround_warns_on_type_drift( + tmp_path: Path, + caplog: pytest.LogCaptureFixture, +) -> None: + """The precise reviewer scenario: nullability + server_default match the + model, only the type is wrong. Pre-family-check, this returned zero + warning (silent JSON-vs-TEXT drift). The family check in + ``_type_equivalent`` must catch this while leaving JSON/JSONB pairs + equivalent so Postgres dialect synonyms don't false-positive.""" + engine = create_async_engine(_url(tmp_path)) + try: + # Reviewer's exact workaround: right nullability/default, wrong type. + await _seed_legacy_without_column(engine) + async with engine.begin() as conn: + await conn.execute(sa.text("ALTER TABLE runs ADD COLUMN token_usage_by_model TEXT NOT NULL DEFAULT '{}'")) + + with caplog.at_level("WARNING", logger="deerflow.persistence.migrations._helpers"): + await bootstrap_schema(engine, backend="sqlite") + + assert await _alembic_version(engine) == HEAD + # No auto-repair: the TEXT column survives unchanged so the operator + # can decide whether to ALTER it themselves. + col = await _runs_column_meta(engine, "token_usage_by_model") + assert col["nullable"] is False + + drift_warnings = [r for r in caplog.records if r.levelname == "WARNING" and r.name == "deerflow.persistence.migrations._helpers" and "safe_add_column" in r.getMessage() and "token_usage_by_model" in r.getMessage()] + assert drift_warnings, "expected safe_add_column to warn about pure type drift (was silent before the family check)" + msg = drift_warnings[0].getMessage() + # The drift entry must explicitly name the type mismatch -- this is + # what was missing before the family check existed. + assert "type actual=" in msg and "desired=" in msg, f"warning missing type drift entry: {msg!r}" + assert "TEXT" in msg and "JSON" in msg, f"warning missing TEXT/JSON in payload: {msg!r}" + # Nullability + server_default match the model -- no other diffs. + assert "nullable" not in msg, f"unexpected nullability drift on matching column: {msg!r}" + assert "server_default" not in msg, f"unexpected server_default drift on matching column: {msg!r}" + finally: + await engine.dispose() + + +# --------------------------------------------------------------------------- +# _type_equivalent unit tests: pin the JSON/JSONB equivalence so Postgres +# dialect synonyms stay quiet, and pin the TEXT/JSON divergence so the +# reviewer's wrong-type scenario keeps firing. +# --------------------------------------------------------------------------- + + +def test_type_equivalent_matches_known_dialect_synonyms() -> None: + from deerflow.persistence.migrations._helpers import _type_equivalent + + # JSON ↔ JSONB (Postgres dialect difference, operationally interchangeable + # for our schema). Both directions, and via raw strings. + assert _type_equivalent(sa.JSON(), "JSONB()") is True + assert _type_equivalent("JSON", "JSONB") is True + assert _type_equivalent("JSONB", "JSON") is True + + +def test_type_equivalent_catches_wholesale_type_mismatch() -> None: + from deerflow.persistence.migrations._helpers import _type_equivalent + + # The reviewer scenario: TEXT NOT NULL DEFAULT '{}' workaround. + assert _type_equivalent("TEXT", "JSON") is False + assert _type_equivalent("TEXT", "JSONB") is False + # Unrelated families also don't accidentally pair up. + assert _type_equivalent("INTEGER", "JSON") is False + + +def test_type_equivalent_ignores_type_parameters() -> None: + """Length / precision differences are out of scope for this helper -- + the goal is wholesale-type drift, not dialect-rendered size defaults.""" + from deerflow.persistence.migrations._helpers import _type_equivalent + + assert _type_equivalent("VARCHAR(255)", "VARCHAR(500)") is True + assert _type_equivalent("NUMERIC(10,2)", "NUMERIC(20,4)") is True + + +def test_type_equivalent_returns_true_on_missing_info() -> None: + """Missing reflected info must not false-positive into a noisy warning.""" + from deerflow.persistence.migrations._helpers import _type_equivalent + + assert _type_equivalent(None, sa.JSON()) is True + assert _type_equivalent(sa.JSON(), None) is True + assert _type_equivalent("", "JSON") is True + + +# --------------------------------------------------------------------------- +# Branch 4: versioned DB +# --------------------------------------------------------------------------- + + +@asyncio_test +async def test_versioned_branch_is_noop_at_head(tmp_path: Path) -> None: + engine = create_async_engine(_url(tmp_path)) + try: + # First bootstrap takes us through the empty branch. + await bootstrap_schema(engine, backend="sqlite") + cols_before = await _runs_columns(engine) + # Second call hits the versioned branch. + await bootstrap_schema(engine, backend="sqlite") + cols_after = await _runs_columns(engine) + assert cols_after == cols_before + assert await _alembic_version(engine) == HEAD + finally: + await engine.dispose() + + +# --------------------------------------------------------------------------- +# Schema-parity guard: legacy-upgraded DB must end up structurally identical +# to a fresh DB on the columns the migration touches. This is the property +# that catches drift between ``Base.metadata`` and ``0002``'s DDL -- exactly +# the failure mode of the original #3682 bug, just at a different layer. +# --------------------------------------------------------------------------- + + +@asyncio_test +async def test_token_usage_column_parity_between_fresh_and_upgraded(tmp_path: Path) -> None: + fresh = create_async_engine(_url(tmp_path, "fresh.db")) + upgraded = create_async_engine(_url(tmp_path, "upgraded.db")) + try: + # Fresh DB -> empty branch -> create_all + await bootstrap_schema(fresh, backend="sqlite") + fresh_col = await _runs_column_meta(fresh, "token_usage_by_model") + + # Legacy DB -> stamp baseline + 0002 upgrade + await _seed_legacy_without_column(upgraded) + await bootstrap_schema(upgraded, backend="sqlite") + upgraded_col = await _runs_column_meta(upgraded, "token_usage_by_model") + + # Pin the contract: the column must have the same nullability AND + # server_default after either bootstrap path. If 0002 ever drifts + # from the model's ``Mapped[dict] = mapped_column(JSON, default=dict, + # server_default=text("'{}'"))`` (i.e. ``nullable=False`` plus the + # ``'{}'`` DB-side default), this fires. + assert fresh_col["nullable"] == upgraded_col["nullable"], f"nullability drift: fresh={fresh_col['nullable']} upgraded={upgraded_col['nullable']}" + # The model declares Mapped[dict] (non-optional) -> NOT NULL. + assert fresh_col["nullable"] is False + assert upgraded_col["nullable"] is False + # Normalize through the same helper the drift warning uses so dialect + # quirks (outer parens, ``::cast``) do not cause false negatives. + assert _normalize_default(fresh_col.get("default")) == _normalize_default(upgraded_col.get("default")), f"server_default drift: fresh={fresh_col.get('default')!r} upgraded={upgraded_col.get('default')!r}" + finally: + await fresh.dispose() + await upgraded.dispose() + + +# --------------------------------------------------------------------------- +# Full schema parity: ``Base.metadata.create_all`` and ``alembic upgrade +# base->head`` MUST produce structurally identical schemas. Both are +# independent sources of the same schema in this codebase -- fresh DBs are +# provisioned by the former (empty branch), historical/upgraded DBs by the +# latter (versioned branch and the alembic tail of the legacy branch). If +# they diverge, two users running the same app version end up with different +# DB structures: exactly the cross-deployment drift this PR exists to kill. +# +# The check is intentionally scoped to columns × (nullable, server_default) +# instead of full type/index/FK reflection. Those are the two highest-signal +# attributes for the drift modes seen so far (#3682 was a nullability +# mismatch; review todo #6 was a server_default mismatch). Type, index, and +# FK reflection differ enough across dialects to require careful +# normalization helpers that aren't worth introducing for this PR's scope; +# see review todo #7 for the wider plan. +# --------------------------------------------------------------------------- + + +def _reflect_columns_sync(sync_conn) -> dict[str, dict[str, dict]]: + insp = sa.inspect(sync_conn) + out: dict[str, dict[str, dict]] = {} + for table in insp.get_table_names(): + # ``alembic_version`` is alembic's own bookkeeping table, not part of + # our schema -- one path creates it (upgrade) and the other doesn't + # (create_all), so comparing it would produce a guaranteed false + # positive every run. + if table == "alembic_version": + continue + out[table] = {c["name"]: c for c in insp.get_columns(table)} + return out + + +async def _reflect_columns(engine) -> dict[str, dict[str, dict]]: + async with engine.connect() as conn: + return await conn.run_sync(_reflect_columns_sync) + + +@asyncio_test +async def test_create_all_and_alembic_upgrade_produce_same_schema(tmp_path: Path) -> None: + fresh = create_async_engine(_url(tmp_path, "fresh.db")) + upgraded = create_async_engine(_url(tmp_path, "upgraded.db")) + try: + # Path A: ``Base.metadata.create_all`` -- the empty-branch code path. + async with fresh.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + # Path B: pure alembic ``upgrade base->head``. Note we deliberately + # bypass ``bootstrap_schema`` on this side -- its empty branch uses + # ``create_all``, not the alembic chain -- to exercise the path a + # versioned-DB upgrade actually takes. + cfg = _get_alembic_config(upgraded) + await asyncio.to_thread(_upgrade, cfg, "head") + + fresh_tables = await _reflect_columns(fresh) + upgraded_tables = await _reflect_columns(upgraded) + + # Same set of tables. A mismatch here means either ``Base.metadata`` + # has gained/lost a table without a matching revision, or a revision + # creates/drops a table without a matching model change. + assert set(fresh_tables) == set(upgraded_tables), f"table-set drift between create_all and alembic upgrade: only-in-create_all={set(fresh_tables) - set(upgraded_tables)} only-in-alembic={set(upgraded_tables) - set(fresh_tables)}" + + for table in sorted(fresh_tables): + fresh_cols = fresh_tables[table] + upgraded_cols = upgraded_tables[table] + assert set(fresh_cols) == set(upgraded_cols), f"{table}: column-set drift only-in-create_all={set(fresh_cols) - set(upgraded_cols)} only-in-alembic={set(upgraded_cols) - set(fresh_cols)}" + for col_name in sorted(fresh_cols): + f_col = fresh_cols[col_name] + u_col = upgraded_cols[col_name] + assert f_col["nullable"] == u_col["nullable"], f"{table}.{col_name}: nullable drift create_all={f_col['nullable']} alembic={u_col['nullable']}" + # Normalize through ``_normalize_default`` to absorb the + # dialect-rendering quirks (outer parens, ``::cast``) that + # would otherwise cause false positives. + f_default = _normalize_default(f_col.get("default")) + u_default = _normalize_default(u_col.get("default")) + assert f_default == u_default, f"{table}.{col_name}: server_default drift create_all={f_col.get('default')!r} alembic={u_col.get('default')!r}" + finally: + await fresh.dispose() + await upgraded.dispose() + + +# --------------------------------------------------------------------------- +# Baseline-table-restriction guards. The legacy branch's backfill must +# create *only* the baseline-era tables, not the full ``Base.metadata``. +# Otherwise it would pre-empt a future ``op.create_table`` revision for a +# newly-added model (the revision would crash with ``relation already +# exists``). Two tests cover this: +# +# 1. ``_BASELINE_TABLE_NAMES`` is pinned against what ``0001_baseline`` +# actually creates -- editing 0001 without updating the constant fires +# here, forcing the developer to keep the two in sync. +# 2. Regression for the leak itself: a phantom table outside the constant +# must NOT be created by the backfill helper. +# --------------------------------------------------------------------------- + + +@asyncio_test +async def test_baseline_table_names_constant_matches_0001(tmp_path: Path) -> None: + engine = create_async_engine(_url(tmp_path)) + try: + cfg = _get_alembic_config(engine) + # Run only up to baseline (not head) and reflect what it produced. + await asyncio.to_thread(_upgrade, cfg, BASELINE) + + async with engine.connect() as conn: + reflected = await conn.run_sync(lambda c: set(sa.inspect(c).get_table_names())) + # ``alembic_version`` is alembic's bookkeeping table, not part of + # our schema -- the constant is about DeerFlow-owned baseline tables. + reflected.discard("alembic_version") + + assert reflected == _BASELINE_TABLE_NAMES, f"_BASELINE_TABLE_NAMES drifted from 0001_baseline.upgrade()'s output: only-in-0001={sorted(reflected - _BASELINE_TABLE_NAMES)} only-in-constant={sorted(_BASELINE_TABLE_NAMES - reflected)}" + finally: + await engine.dispose() + + +@asyncio_test +async def test_legacy_backfill_skips_non_baseline_tables(tmp_path: Path) -> None: + """Regression: legacy backfill must not create tables outside the baseline + set, because a later ``op.create_table`` revision for the same name would + fail. We synthesise a phantom table on ``Base.metadata`` (modelling a + future model addition), run the backfill helper, and assert the phantom + is absent from the resulting DB. + """ + phantom_name = "phantom_future_table_for_test" + phantom = sa.Table( + phantom_name, + Base.metadata, + sa.Column("id", sa.Integer, primary_key=True), + ) + try: + engine = create_async_engine(_url(tmp_path)) + try: + async with engine.begin() as conn: + await conn.run_sync(_run_baseline_create_all_sync) + + async with engine.connect() as conn: + tables = await conn.run_sync(lambda c: set(sa.inspect(c).get_table_names())) + + assert phantom_name not in tables, f"legacy backfill leaked {phantom_name!r}; a future ``op.create_table({phantom_name!r})`` revision would now collide" + # Sanity: baseline tables ARE created by the backfill helper. + assert "runs" in tables + assert "channel_connections" in tables + finally: + await engine.dispose() + finally: + Base.metadata.remove(phantom) + + +# --------------------------------------------------------------------------- +# _decide_state unit tests (pure function, no DB needed) +# --------------------------------------------------------------------------- + + +class TestDecideState: + def test_empty(self): + assert _decide_state({"has_alembic_version": False, "has_deerflow_tables": False}) == "empty" + + def test_empty_with_unrelated_tables(self): + # LangGraph checkpointer tables present but DeerFlow has nothing yet. + # ``has_deerflow_tables`` is derived from the metadata intersection in + # production, so the only thing the decision function needs is the + # bool itself. + assert _decide_state({"has_alembic_version": False, "has_deerflow_tables": False}) == "empty" + + def test_legacy(self): + assert _decide_state({"has_alembic_version": False, "has_deerflow_tables": True}) == "legacy" + + def test_versioned(self): + assert _decide_state({"has_alembic_version": True, "has_deerflow_tables": True}) == "versioned" + + def test_versioned_takes_precedence_over_empty(self): + # Pathological: alembic_version row exists but no managed tables yet + # (e.g. someone restored only the alembic_version table from backup). + # We still go versioned -> upgrade head, which is the right thing: + # alembic will run every revision from base. + assert _decide_state({"has_alembic_version": True, "has_deerflow_tables": False}) == "versioned" + + +# --------------------------------------------------------------------------- +# Sanity: head revision is the one this module expects +# --------------------------------------------------------------------------- + + +def test_head_revision_is_token_usage_revision() -> None: + assert _get_head_revision() == HEAD + + +def test_baseline_revision_id_is_known() -> None: + """Detect a baseline rename: the bootstrap code hardcodes ``0001_baseline`` + as the stamp target for the legacy branch, so a rename would silently + break that branch unless caught here.""" + from pathlib import Path # noqa: PLC0415 + + from alembic.config import Config # noqa: PLC0415 + from alembic.script import ScriptDirectory # noqa: PLC0415 + + migrations_dir = Path(__file__).resolve().parents[1] / "packages/harness/deerflow/persistence/migrations" + cfg = Config() + cfg.set_main_option("script_location", str(migrations_dir)) + script = ScriptDirectory.from_config(cfg) + all_ids = {rev.revision for rev in script.walk_revisions()} + assert BASELINE in all_ids, f"baseline revision id {BASELINE!r} not found in {all_ids}" diff --git a/backend/tests/test_persistence_bootstrap_concurrency.py b/backend/tests/test_persistence_bootstrap_concurrency.py new file mode 100644 index 000000000..e3c0a3b64 --- /dev/null +++ b/backend/tests/test_persistence_bootstrap_concurrency.py @@ -0,0 +1,150 @@ +"""Concurrency safety tests for ``bootstrap_schema``. + +The contract: N concurrent callers against the same DB always converge to +``alembic_version == head`` without exceptions and without duplicate schema +mutations. + +We model concurrency at the *async-task* level here (multiple coroutines +inside one process). SQLite is single-node by deployment, so within-process +serialisation -- which is what the per-engine ``_SQLITE_LOCKS`` entry +provides -- is the realistic boundary. Cross-process serialisation falls +through to SQLite's own write lock + ``PRAGMA busy_timeout`` plus the +idempotent revision helpers. +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path + +import pytest +import sqlalchemy as sa +from sqlalchemy.ext.asyncio import create_async_engine + +import deerflow.persistence.models # noqa: F401 +from deerflow.persistence import bootstrap as bootstrap_mod +from deerflow.persistence.bootstrap import bootstrap_schema + +pytestmark = pytest.mark.asyncio + + +HEAD = "0002_runs_token_usage" + + +def _url(tmp_path: Path) -> str: + return f"sqlite+aiosqlite:///{(tmp_path / 'concurrent.db').as_posix()}" + + +async def _alembic_version(engine) -> str | None: + async with engine.connect() as conn: + row = await conn.execute(sa.text("SELECT version_num FROM alembic_version")) + return row.scalar() + + +async def _runs_columns(engine) -> set[str]: + async with engine.connect() as conn: + return await conn.run_sync(lambda c: {col["name"] for col in sa.inspect(c).get_columns("runs")}) + + +async def test_two_concurrent_bootstrap_callers_converge(tmp_path: Path) -> None: + engine = create_async_engine(_url(tmp_path)) + try: + await asyncio.gather( + bootstrap_schema(engine, backend="sqlite"), + bootstrap_schema(engine, backend="sqlite"), + ) + assert await _alembic_version(engine) == HEAD + assert "token_usage_by_model" in await _runs_columns(engine) + finally: + await engine.dispose() + + +async def test_five_concurrent_bootstrap_callers_converge(tmp_path: Path) -> None: + engine = create_async_engine(_url(tmp_path)) + try: + await asyncio.gather(*(bootstrap_schema(engine, backend="sqlite") for _ in range(5))) + assert await _alembic_version(engine) == HEAD + finally: + await engine.dispose() + + +async def test_cancelled_caller_does_not_block_others(tmp_path: Path) -> None: + """Cancelling one task mid-bootstrap must not strand the lock or the DB. + + After the cancel, a subsequent ``bootstrap_schema`` call must still reach + head. + """ + engine = create_async_engine(_url(tmp_path)) + try: + task = asyncio.create_task(bootstrap_schema(engine, backend="sqlite")) + # Give the event loop a turn so the task can start; then cancel. + await asyncio.sleep(0) + task.cancel() + # Cancelled task may have raced past the lock; swallow either outcome. + try: + await task + except (asyncio.CancelledError, Exception): # noqa: BLE001 + pass + + # Lock must be free for the next caller. + await bootstrap_schema(engine, backend="sqlite") + assert await _alembic_version(engine) == HEAD + finally: + await engine.dispose() + + +async def test_late_caller_after_head_is_noop(monkeypatch, tmp_path: Path) -> None: + """When the first caller leaves the DB at head, the second observes + 'versioned' and skips create_all / stamp -- it only runs upgrade head, + which is alembic-no-op. + + We use a monkeypatched ``_upgrade`` counter to assert the second caller's + upgrade ran but did no real work (no new revision applied). + """ + engine = create_async_engine(_url(tmp_path)) + try: + # First caller: empty branch. + await bootstrap_schema(engine, backend="sqlite") + first_version = await _alembic_version(engine) + assert first_version == HEAD + + upgrade_calls: list[str] = [] + original_upgrade = bootstrap_mod._upgrade + + def counting_upgrade(cfg, rev: str) -> None: + upgrade_calls.append(rev) + original_upgrade(cfg, rev) + + monkeypatch.setattr(bootstrap_mod, "_upgrade", counting_upgrade) + + # Second caller: versioned branch -> calls _upgrade('head'). + await bootstrap_schema(engine, backend="sqlite") + assert upgrade_calls == ["head"] + assert await _alembic_version(engine) == HEAD + finally: + await engine.dispose() + + +async def test_slow_upgrade_does_not_corrupt_concurrent_state(monkeypatch, tmp_path: Path) -> None: + """Inject a delay into the upgrade path; concurrent callers must still + converge to head with no exceptions.""" + engine = create_async_engine(_url(tmp_path)) + try: + original_upgrade = bootstrap_mod._upgrade + + def slow_upgrade(cfg, rev: str) -> None: + import time # noqa: PLC0415 + + time.sleep(0.2) + original_upgrade(cfg, rev) + + monkeypatch.setattr(bootstrap_mod, "_upgrade", slow_upgrade) + + await asyncio.gather( + bootstrap_schema(engine, backend="sqlite"), + bootstrap_schema(engine, backend="sqlite"), + bootstrap_schema(engine, backend="sqlite"), + ) + assert await _alembic_version(engine) == HEAD + finally: + await engine.dispose() diff --git a/backend/tests/test_persistence_bootstrap_pg_lock.py b/backend/tests/test_persistence_bootstrap_pg_lock.py new file mode 100644 index 000000000..5755d3b2e --- /dev/null +++ b/backend/tests/test_persistence_bootstrap_pg_lock.py @@ -0,0 +1,98 @@ +"""Regression test for the Postgres bootstrap advisory-lock protection. + +Managed Postgres (RDS, Cloud SQL, Supabase) defaults +``idle_in_transaction_session_timeout`` to 1-10 minutes. If the lock-holding +connection sits idle while ``asyncio.to_thread(_upgrade, ...)`` runs alembic +on a different pooled connection longer than that, the host kills the idle +session and the advisory lock is **silently released** -- defeating the +cross-process mutex. ``_postgres_lock`` issues +``SET LOCAL idle_in_transaction_session_timeout = 0`` immediately on the +lock-holding connection to neutralise that kill for the lifetime of the +transaction. + +This test pins: + +1. The ``SET LOCAL`` is emitted at all (no silent regression). +2. It runs **before** ``pg_advisory_lock`` -- otherwise a slow lock acquire + on a heavily-contended cluster would itself be vulnerable. +3. The ``pg_advisory_unlock`` still fires on the way out (the new SQL must + not break the release path). + +We mock the engine instead of standing up a real Postgres because the only +behaviour worth pinning here is the SQL execution order; the timeout's +runtime effect is Postgres's contract, not ours. +""" + +from __future__ import annotations + +import pytest + +from deerflow.persistence import bootstrap as bootstrap_mod + + +class _FakeAsyncConn: + """Async-context-manager stand-in for SQLAlchemy's ``AsyncConnection``. + + Records every ``execute(stmt, params)`` so the test can assert SQL order. + """ + + def __init__(self) -> None: + self.executed: list[tuple[str, dict | None]] = [] + + async def execute(self, stmt, params=None): + self.executed.append((str(stmt), params)) + return None + + async def __aenter__(self) -> _FakeAsyncConn: + return self + + async def __aexit__(self, *_exc_info: object) -> None: + return None + + +class _FakeAsyncEngine: + def __init__(self) -> None: + self.conn = _FakeAsyncConn() + + def connect(self) -> _FakeAsyncConn: + return self.conn + + +@pytest.mark.asyncio +async def test_postgres_lock_disables_idle_in_transaction_kill_before_locking() -> None: + engine = _FakeAsyncEngine() + + async with bootstrap_mod._postgres_lock(engine): # type: ignore[arg-type] + pass + + sqls = [stmt for stmt, _ in engine.conn.executed] + + # 1. SET LOCAL fires. + set_local_idx = next( + (i for i, s in enumerate(sqls) if "set local idle_in_transaction_session_timeout" in s.lower()), + None, + ) + assert set_local_idx is not None, f"SET LOCAL never executed; saw: {sqls}" + assert "0" in sqls[set_local_idx], f"SET LOCAL did not target value 0: {sqls[set_local_idx]!r}" + + # 2. SET LOCAL precedes pg_advisory_lock. + lock_idx = next((i for i, s in enumerate(sqls) if "pg_advisory_lock" in s), None) + assert lock_idx is not None, f"pg_advisory_lock never executed; saw: {sqls}" + assert set_local_idx < lock_idx, f"SET LOCAL must run before pg_advisory_lock; got order {sqls}" + + # 3. pg_advisory_unlock still fires on exit. + assert any("pg_advisory_unlock" in s for s in sqls), f"pg_advisory_unlock missing; saw: {sqls}" + + +@pytest.mark.asyncio +async def test_postgres_lock_releases_even_if_body_raises() -> None: + """Defence-in-depth: the SET LOCAL addition must not regress the + existing finally-block contract that releases the lock on body errors.""" + engine = _FakeAsyncEngine() + + with pytest.raises(RuntimeError, match="boom"): + async with bootstrap_mod._postgres_lock(engine): # type: ignore[arg-type] + raise RuntimeError("boom") + + sqls = [stmt for stmt, _ in engine.conn.executed] + assert any("pg_advisory_unlock" in s for s in sqls), f"unlock missing after body error; saw: {sqls}" diff --git a/backend/tests/test_persistence_bootstrap_regression.py b/backend/tests/test_persistence_bootstrap_regression.py new file mode 100644 index 000000000..546311b13 --- /dev/null +++ b/backend/tests/test_persistence_bootstrap_regression.py @@ -0,0 +1,121 @@ +"""Regression test for GitHub issue #3682. + +End-to-end shape: + +1. Hand-build a SQLite DB that mirrors a real pre-#3658 deployment -- the + ``runs`` table is missing the ``token_usage_by_model`` column, mirroring + what every existing user's DB looked like after the upgrade that triggered + the issue. +2. Run ``init_engine`` (the entry point used by the FastAPI Gateway + lifespan), which now routes through ``bootstrap_schema``. +3. Confirm a real ``SELECT`` against the column succeeds, demonstrating the + 500 from the original issue is gone. + +The pre-fix codepath would have raised +``sqlalchemy.exc.OperationalError: no such column: runs.token_usage_by_model`` +on step 3. +""" + +from __future__ import annotations + +import sqlite3 +from pathlib import Path +from uuid import uuid4 + +import pytest +import sqlalchemy as sa + +import deerflow.persistence.models # noqa: F401 -- registers ORM models +from deerflow.persistence.base import Base +from deerflow.persistence.engine import close_engine, get_session_factory, init_engine +from deerflow.persistence.run import RunRepository + +pytestmark = pytest.mark.asyncio + + +def _seed_pre_3658_database(db_path: Path) -> None: + """Build a DB that looks like a pre-PR-#3658 deployment. + + Uses the synchronous ``sqlite3`` driver so the seed is independent of the + async engine under test. + """ + db_path.parent.mkdir(parents=True, exist_ok=True) + + # Easiest way to get the legacy shape exactly right: create_all then + # ALTER away the new column. + sync_url = f"sqlite:///{db_path.as_posix()}" + sync_engine = sa.create_engine(sync_url) + try: + Base.metadata.create_all(sync_engine) + with sync_engine.begin() as conn: + conn.execute(sa.text("ALTER TABLE runs DROP COLUMN token_usage_by_model")) + finally: + sync_engine.dispose() + + +async def test_legacy_database_recovers_token_usage_column(tmp_path: Path) -> None: + db_path = tmp_path / "legacy.db" + _seed_pre_3658_database(db_path) + + # Sanity: confirm we did indeed land in the buggy pre-fix shape before + # init_engine touches the file. + with sqlite3.connect(db_path) as raw: + cols = {row[1] for row in raw.execute("PRAGMA table_info(runs)").fetchall()} + assert "run_id" in cols + assert "token_usage_by_model" not in cols + version_table_count = raw.execute("SELECT count(*) FROM sqlite_master WHERE type='table' AND name='alembic_version'").fetchone()[0] + assert version_table_count == 0 + + # Run the same init_engine path FastAPI lifespan uses on startup. + url = f"sqlite+aiosqlite:///{db_path.as_posix()}" + await init_engine(backend="sqlite", url=url, sqlite_dir=str(tmp_path)) + + try: + # The column must now be present. + with sqlite3.connect(db_path) as raw: + cols = {row[1] for row in raw.execute("PRAGMA table_info(runs)").fetchall()} + assert "token_usage_by_model" in cols + version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone() + assert version_row[0] == "0002_runs_token_usage" + + # And the read path that originally 500'd must now succeed. + sf = get_session_factory() + assert sf is not None + repo = RunRepository(sf) + # No rows yet -- the point is just that the SELECT does not raise + # ``no such column: runs.token_usage_by_model``. + result = await repo.aggregate_tokens_by_thread(thread_id=str(uuid4())) + assert result["total_tokens"] == 0 + assert result["by_model"] == {} + finally: + await close_engine() + + +async def test_legacy_database_with_manual_alter_still_bootstraps(tmp_path: Path) -> None: + """User-side workaround scenario: someone already applied the manual + ``ALTER TABLE runs ADD COLUMN token_usage_by_model JSON`` from the issue + write-up. The hybrid bootstrap must just stamp head, not double-add the + column, and not error. + """ + db_path = tmp_path / "manual_altered.db" + db_path.parent.mkdir(parents=True, exist_ok=True) + + sync_engine = sa.create_engine(f"sqlite:///{db_path.as_posix()}") + try: + Base.metadata.create_all(sync_engine) + # Don't strip the column -- this is the "user already ran the + # workaround" case where create_all already produced it. + finally: + sync_engine.dispose() + + url = f"sqlite+aiosqlite:///{db_path.as_posix()}" + await init_engine(backend="sqlite", url=url, sqlite_dir=str(tmp_path)) + try: + with sqlite3.connect(db_path) as raw: + cols = [row[1] for row in raw.execute("PRAGMA table_info(runs)").fetchall()] + # No duplicate column -- list, not set, to catch dupes. + assert cols.count("token_usage_by_model") == 1 + version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone() + assert version_row[0] == "0002_runs_token_usage" + finally: + await close_engine() diff --git a/backend/tests/test_persistence_bootstrap_sqlite_lock.py b/backend/tests/test_persistence_bootstrap_sqlite_lock.py new file mode 100644 index 000000000..ebe5b8cfd --- /dev/null +++ b/backend/tests/test_persistence_bootstrap_sqlite_lock.py @@ -0,0 +1,116 @@ +"""Regression tests for the per-engine SQLite bootstrap lock cache. + +The cache (``deerflow.persistence.bootstrap._SQLITE_LOCKS``) maps an engine +to the ``asyncio.Lock`` that serialises its in-process bootstrap. It is keyed +by the engine object itself via ``WeakKeyDictionary`` -- not ``id(engine)`` -- +to avoid two failure modes that are silent in production (one long-lived +engine) but real in pytest (one fresh engine per test): + +1. **CPython id reuse.** After an engine is garbage-collected its memory + address can be reused by a new engine. An ``id``-keyed cache would hand + the new engine the dead engine's ``Lock``. That lock was bound to the + dead engine's event loop at first ``async with``; pytest gives each async + test its own loop, so reusing it raises ``RuntimeError: ... bound to a + different event loop``. +2. **Unbounded growth.** An ``id``-keyed cache never drops entries because + nothing notifies it when the engine dies. With ``WeakKeyDictionary`` the + entry disappears as soon as the engine is collected. + +These tests do not open any DB connection -- they exercise the cache helper +directly so they can run without an event loop and without aiosqlite warnings +about unclosed engines. +""" + +from __future__ import annotations + +import gc +import weakref + +import pytest +from sqlalchemy.ext.asyncio import create_async_engine + +from deerflow.persistence import bootstrap as bootstrap_mod +from deerflow.persistence.bootstrap import _get_sqlite_local_lock + + +def _make_engine(): + return create_async_engine("sqlite+aiosqlite:///:memory:") + + +def test_cache_is_weak_key_dictionary() -> None: + """Pin the cache type so a refactor cannot silently revert to a plain + dict (which would reintroduce the id-reuse bug).""" + assert isinstance(bootstrap_mod._SQLITE_LOCKS, weakref.WeakKeyDictionary) + + +def test_same_engine_returns_same_lock() -> None: + engine = _make_engine() + assert _get_sqlite_local_lock(engine) is _get_sqlite_local_lock(engine) + + +def test_distinct_engines_get_distinct_locks() -> None: + """Two live engines must not share a lock -- otherwise unrelated + bootstraps would serialise against each other.""" + engine_a = _make_engine() + engine_b = _make_engine() + assert _get_sqlite_local_lock(engine_a) is not _get_sqlite_local_lock(engine_b) + + +def test_entry_drops_when_engine_is_garbage_collected() -> None: + """The cache must not pin the engine alive. + + This is the structural guarantee behind the id-reuse fix: when the engine + is collected, its lock entry goes with it, so a future engine landing on + the same address cannot inherit a stale, loop-bound lock. + """ + engine = _make_engine() + _get_sqlite_local_lock(engine) + assert engine in bootstrap_mod._SQLITE_LOCKS + + engine_ref = weakref.ref(engine) + del engine + gc.collect() + + assert engine_ref() is None, "engine should be collectible -- cache must not hold a strong ref" + # WeakKeyDictionary may defer removal until the next access; touch it. + assert all(ref() is not None for ref in bootstrap_mod._SQLITE_LOCKS.keyrefs()) + + +@pytest.mark.asyncio +async def test_fresh_engine_gets_lock_usable_on_current_loop() -> None: + """End-to-end guard for the pytest pattern: a brand-new engine in a + brand-new event loop must receive a lock that ``async with`` accepts. + + This is the behaviour an ``id``-keyed cache could break if the new engine + landed on a previously-used address -- it would return a lock bound to a + dead loop and raise ``RuntimeError: ... bound to a different event loop``. + """ + engine = _make_engine() + try: + lock = _get_sqlite_local_lock(engine) + async with lock: + pass + # Re-entrant acquire on the same loop must also succeed. + async with lock: + pass + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_cache_does_not_grow_across_disposed_engines() -> None: + """Create + dispose + drop many engines and assert the cache stays bounded. + + Without ``WeakKeyDictionary`` this loop would leak one entry per engine. + """ + initial = len(bootstrap_mod._SQLITE_LOCKS) + for _ in range(20): + engine = _make_engine() + _get_sqlite_local_lock(engine) + await engine.dispose() + del engine + gc.collect() + # Touch the dict so WeakKeyDictionary clears any deferred removals. + _ = list(bootstrap_mod._SQLITE_LOCKS.items()) + # Allow a small slack for any engine that is still pinned by a frame. + assert len(bootstrap_mod._SQLITE_LOCKS) - initial <= 1 diff --git a/backend/tests/test_persistence_bootstrap_url.py b/backend/tests/test_persistence_bootstrap_url.py new file mode 100644 index 000000000..d29799c71 --- /dev/null +++ b/backend/tests/test_persistence_bootstrap_url.py @@ -0,0 +1,71 @@ +"""Tests for the Postgres URL / ConfigParser pitfalls in ``bootstrap``. + +Two failure modes the ``_alembic_safe_url`` helper exists to prevent: + +1. ``str(engine.url)`` (and the default ``URL.render_as_string()``) masks the + password as ``***``. The live engine would still work because it carries + the password in-memory, but alembic ``stamp`` / ``upgrade`` (which open + their own connection from the URL we pass in) would authenticate with + garbage and fail at runtime. +2. ``alembic.config.Config.set_main_option`` forwards to ``ConfigParser.set``, + which performs ``%(name)s``-style interpolation on the value. A URL-encoded + password containing ``%`` (e.g. ``p%40ss`` for ``p@ss``) raises + ``InterpolationSyntaxError``. Every literal ``%`` must be doubled. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +from sqlalchemy.engine.url import make_url + +from deerflow.persistence.bootstrap import _alembic_safe_url, _escape_url_for_alembic, _get_alembic_config + + +def _fake_engine(url: str) -> SimpleNamespace: + """Build a minimal stand-in for ``AsyncEngine`` so we don't need a real + driver (e.g. asyncpg) installed just to exercise the URL path.""" + return SimpleNamespace(url=make_url(url)) + + +def test_safe_url_preserves_password_for_postgres() -> None: + engine = _fake_engine("postgresql://alice:s3cret@db.example.com/app") + safe = _alembic_safe_url(engine) + assert "s3cret" in safe, "password got masked: alembic would auth with garbage" + assert "***" not in safe + + +def test_safe_url_escapes_percent_for_configparser() -> None: + # URL-encoded ``@`` in password -> raw ``%40`` in URL -> ConfigParser + # would treat it as an interpolation marker. + engine = _fake_engine("postgresql://alice:p%40ss@db.example.com/app") + safe = _alembic_safe_url(engine) + assert "p%%40ss" in safe, f"percent not doubled, ConfigParser will fail: {safe}" + + +def test_alembic_config_accepts_url_with_percent_and_round_trips() -> None: + # The whole point: build_config should not raise, and the URL alembic + # reads back should match the original (single ``%``, real password). + original = "postgresql://alice:p%40ss@db.example.com/app" + engine = _fake_engine(original) + cfg = _get_alembic_config(engine) + roundtrip = cfg.get_main_option("sqlalchemy.url") + assert roundtrip == original, f"alembic sees a different URL than we set: {roundtrip}" + + +def test_sqlite_url_does_not_double_percent_unnecessarily() -> None: + # No percent in the URL -> no escaping needed -> output equals input. + engine = _fake_engine("sqlite+aiosqlite:///tmp/db.sqlite") + safe = _alembic_safe_url(engine) + assert safe == "sqlite+aiosqlite:///tmp/db.sqlite" + + +def test_escape_url_for_alembic_doubles_only_percent_signs() -> None: + # Shared helper used by both ``bootstrap._alembic_safe_url`` and + # ``scripts/_autogen_revision._alembic_config`` -- pins the round-trip + # rule so any future URL/ConfigParser corner case is fixed in one place. + assert _escape_url_for_alembic("postgresql://a:p%40ss@h/d") == "postgresql://a:p%%40ss@h/d" + assert _escape_url_for_alembic("sqlite:///x.db") == "sqlite:///x.db" + # Idempotency is intentionally NOT a property -- doubling is one-way; + # callers must escape exactly once on the way into set_main_option. + assert _escape_url_for_alembic("a%%b") == "a%%%%b" diff --git a/backend/tests/test_persistence_migrations_env.py b/backend/tests/test_persistence_migrations_env.py new file mode 100644 index 000000000..4ca4cd7e1 --- /dev/null +++ b/backend/tests/test_persistence_migrations_env.py @@ -0,0 +1,91 @@ +"""Tests for the ``include_object`` filter used by ``migrations/env.py``. + +LangGraph checkpointer tables (``checkpoints`` and friends) live alongside +DeerFlow's own tables in the same database. Alembic must NEVER emit DDL for +them or a future ``alembic revision --autogenerate`` would propose +``drop_table('checkpoints')`` whenever LangGraph's tables are reflected from +a live DB. + +The filter is the only line of defence between an honest autogenerate run +and a destructive revision. It lives in ``_env_filters.py`` so it can be unit +tested without alembic's import-time machinery. +""" + +from __future__ import annotations + +import sqlalchemy as sa + +from deerflow.persistence.migrations._env_filters import ( + LANGGRAPH_OWNED_TABLES, + include_object, +) + + +def _table(name: str) -> sa.Table: + return sa.Table(name, sa.MetaData()) + + +def test_filter_excludes_langgraph_checkpoint_tables() -> None: + for owned in ( + "checkpoints", + "checkpoint_blobs", + "checkpoint_writes", + "checkpoint_migrations", + ): + assert include_object(_table(owned), owned, "table", True, None) is False + + +def test_filter_includes_deerflow_tables() -> None: + for owned in ("runs", "threads_meta", "feedback", "users", "channel_connections"): + assert include_object(_table(owned), owned, "table", True, None) is True + + +def test_filter_excludes_indexes_on_langgraph_tables() -> None: + # An Index whose parent table is LangGraph-owned must also be filtered out; + # otherwise autogenerate would emit drop_index against tables alembic does + # not own. + md = sa.MetaData() + parent = sa.Table("checkpoints", md, sa.Column("id", sa.Integer, primary_key=True)) + idx = sa.Index("ix_checkpoints_anything", parent.c.id) + assert include_object(idx, idx.name, "index", True, None) is False + + +def test_filter_includes_indexes_on_deerflow_tables() -> None: + md = sa.MetaData() + parent = sa.Table("runs", md, sa.Column("run_id", sa.String, primary_key=True)) + idx = sa.Index("ix_runs_something", parent.c.run_id) + assert include_object(idx, idx.name, "index", True, None) is True + + +def test_langgraph_owned_tables_set_is_complete() -> None: + # Pin the explicit set so an inadvertent removal -- e.g. someone simplifying + # the filter -- requires a test diff that surfaces the change. + assert LANGGRAPH_OWNED_TABLES == frozenset( + { + "checkpoints", + "checkpoint_blobs", + "checkpoint_writes", + "checkpoint_migrations", + } + ) + + +def test_env_module_wires_busy_timeout_for_sqlite() -> None: + """Regression for the cross-process bootstrap pitfall: alembic spawns its + own engine inside ``env.py::run_migrations_online`` and that engine does + NOT inherit PRAGMAs from the production engine. Without an event listener + here, its connections would use the default 5s busy_timeout and racy + multi-process bootstrap would fail with ``database is locked`` instead of + waiting for the file lock. + + We check the source rather than execute env.py (which would try to drive + alembic on import) so this test stays a pure parity check. + """ + from pathlib import Path # noqa: PLC0415 + + env_path = Path(__file__).resolve().parents[1] / "packages/harness/deerflow/persistence/migrations/env.py" + src = env_path.read_text(encoding="utf-8") + assert "PRAGMA busy_timeout=30000" in src or "PRAGMA busy_timeout = 30000" in src, ( + "env.py must set busy_timeout on its alembic-spawned engine; without it, cross-process bootstrap on SQLite fails fast instead of waiting for the file lock" + ) + assert 'listens_for(connectable.sync_engine, "connect")' in src, "busy_timeout must be wired via an event listener so EVERY connection alembic opens gets the PRAGMA, not just one initial probe"