mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-01 19:06:01 +00:00
feat(persistence): wire alembic migrations, bootstrap schema on startup (#3706)
* feat(persistence): wire alembic migrations + bootstrap schema on startup Closes #3682. Pre-#3658 DBs lack the `runs.token_usage_by_model` column because alembic was never wired up — startup only ran `create_all`, which never ALTERs existing tables. Adds a hybrid bootstrap in FastAPI lifespan (replaces bare `create_all`): - empty DB → create_all + stamp head - legacy DB → stamp 0001_baseline + upgrade head - versioned DB → upgrade head Concurrency: Postgres `pg_advisory_lock` (cross-process); SQLite per-engine `asyncio.Lock` + 30s `PRAGMA busy_timeout` on both prod and alembic engines. Column revisions use `safe_add_column` / `safe_drop_column` idempotent helpers as fallback. Other bits: - 0001 baseline (chain root) + 0002 add `runs.token_usage_by_model` - `include_object` filter so alembic ignores LangGraph checkpointer tables - `make migrate-rev MSG="..."` for authoring new revisions (no migrate/stamp targets — startup is the only execution path) - Tests: three-branch decision, concurrency, #3682 regression, env filter, blocking-IO gate anchor - CLAUDE.md: new "Schema migrations" section * fix(style): fix lint error * perf(persistence): address review feedback on alembic bootstrap Behavioural fixes - _SQLITE_LOCKS now keyed via WeakKeyDictionary so id-reuse after GC cannot return a stale, loop-bound lock and the cache cannot leak one entry per disposed engine. - safe_add_column compares nullable / server_default against the desired column when the name already exists and emits a warning on drift, surfacing manual-ALTER workarounds instead of silently no-op'ing. - _postgres_lock issues SET LOCAL idle_in_transaction_session_timeout=0 before pg_advisory_lock, so managed Postgres cannot kill the idle lock-holding session mid-upgrade and silently release the advisory lock. - legacy branch now backfills missing baseline tables via a restricted create_all (Base.metadata.create_all scoped to _BASELINE_TABLE_NAMES). Restores pre-#1930 upgraders whose channel_* tables were never provisioned, without pre-empting future create_table revisions for newly-added models. Schema parity - runs.token_usage_by_model gains server_default=text("'{}'") in both the ORM model and the 0001_baseline create_table, matching what 0002 adds via ALTER. create_all and alembic-upgrade paths now produce identical column definitions. - New parity test compares Base.metadata.create_all output against a pure alembic upgrade base->head, asserting column-set, nullable, and server_default agree across all tables (normalized through the same helper safe_add_column's drift check uses). Guards - test_baseline_table_names_constant_matches_0001 pins _BASELINE_TABLE_NAMES to 0001_baseline.upgrade()'s actual output -- the constant cannot drift silently when someone edits 0001. - test_legacy_backfill_skips_non_baseline_tables verifies the restricted backfill does not create a phantom table on Base.metadata, modelling a future revision that would otherwise collide on op.create_table. Doc residuals - Three-branch decision table is now consistent across bootstrap.py top docstring, engine.py comment, test module docstring, and CLAUDE.md. - Stale test anchor in blocking_io/test_persistence_engine_sqlite.py docstring now points at the real file. * fix(style): fix lint error * fix(persistence): close drift detection holes - _check_column_drift compares column type via a family equivalence allowlist ({JSON, JSONB}). Catches the wrong-type workaround `TEXT NOT NULL DEFAULT '{}'` that previously slipped through silently, while keeping Postgres JSON/JSONB dialect reflection quiet. Reflected and desired type are also echoed in every drift warning's payload for operator triage. - Extract _escape_url_for_alembic so bootstrap._alembic_safe_url and scripts/_autogen_revision share the ConfigParser % escape rule instead of duplicating it. - backend/README.md: add `make migrate-rev MSG=...` to Commands and a Schema Migrations section per the repo's README/CLAUDE.md sync policy. - test_base_to_dict.py: scope the test ORM class to an isolated MetaData so the create_all-vs-alembic parity test (added in the previous commit) is not polluted by the phantom table on the full pytest session.
This commit is contained in:
AnoobFeng
GitHub
parent
346164d740
commit
debb0fd161
@ -95,6 +95,7 @@ make test # Run all backend tests
|
|||||||
make test-blocking-io # Run strict Blockbuster runtime gate on tests/blocking_io/
|
make test-blocking-io # Run strict Blockbuster runtime gate on tests/blocking_io/
|
||||||
make lint # Lint with ruff
|
make lint # Lint with ruff
|
||||||
make format # Format code 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/`,
|
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_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
|
- `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/`)
|
### Tracing System (`packages/harness/deerflow/tracing/`)
|
||||||
|
|
||||||
LangSmith and Langfuse are both supported. The wiring lives in two layers:
|
LangSmith and Langfuse are both supported. The wiring lives in two layers:
|
||||||
|
|||||||
@ -22,3 +22,14 @@ format:
|
|||||||
|
|
||||||
detect-blocking-io:
|
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
|
@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)"
|
||||||
|
|||||||
@ -364,8 +364,35 @@ make gateway # Run Gateway API without reload (port 8001)
|
|||||||
make lint # Run linter (ruff)
|
make lint # Run linter (ruff)
|
||||||
make format # Format code (ruff)
|
make format # Format code (ruff)
|
||||||
make detect-blocking-io # Inventory blocking IO that may block the backend event loop
|
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
|
### Code Style
|
||||||
|
|
||||||
- **Linter/Formatter**: `ruff`
|
- **Linter/Formatter**: `ruff`
|
||||||
|
|||||||
452
backend/packages/harness/deerflow/persistence/bootstrap.py
Normal file
452
backend/packages/harness/deerflow/persistence/bootstrap.py
Normal file
@ -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)
|
||||||
@ -112,11 +112,14 @@ async def init_engine(
|
|||||||
# SQLite deployment (TC-UPG-06 in AUTH_TEST_PLAN.md). The companion
|
# SQLite deployment (TC-UPG-06 in AUTH_TEST_PLAN.md). The companion
|
||||||
# ``synchronous=NORMAL`` is the safe-and-fast pairing — fsync only
|
# ``synchronous=NORMAL`` is the safe-and-fast pairing — fsync only
|
||||||
# at WAL checkpoint boundaries instead of every commit.
|
# at WAL checkpoint boundaries instead of every commit.
|
||||||
# Note: we do not set PRAGMA busy_timeout here — Python's sqlite3
|
# We also widen ``busy_timeout`` to 30s here. Python's sqlite3 driver
|
||||||
# driver already defaults to a 5-second busy timeout (see the
|
# defaults to 5s, which is fine for transient row contention but too
|
||||||
# ``timeout`` kwarg of ``sqlite3.connect``), and aiosqlite /
|
# tight for cross-process bootstrap: the second-N-th Gateway process
|
||||||
# SQLAlchemy's aiosqlite dialect inherit that default. Setting
|
# may need to wait while the first runs ``ALTER TABLE`` /
|
||||||
# it again would be a no-op.
|
# ``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")
|
@event.listens_for(_engine.sync_engine, "connect")
|
||||||
def _enable_sqlite_wal(dbapi_conn, _record): # noqa: ARG001 — SQLAlchemy contract
|
def _enable_sqlite_wal(dbapi_conn, _record): # noqa: ARG001 — SQLAlchemy contract
|
||||||
cursor = dbapi_conn.cursor()
|
cursor = dbapi_conn.cursor()
|
||||||
@ -124,6 +127,7 @@ async def init_engine(
|
|||||||
cursor.execute("PRAGMA journal_mode=WAL;")
|
cursor.execute("PRAGMA journal_mode=WAL;")
|
||||||
cursor.execute("PRAGMA synchronous=NORMAL;")
|
cursor.execute("PRAGMA synchronous=NORMAL;")
|
||||||
cursor.execute("PRAGMA foreign_keys=ON;")
|
cursor.execute("PRAGMA foreign_keys=ON;")
|
||||||
|
cursor.execute("PRAGMA busy_timeout=30000;")
|
||||||
finally:
|
finally:
|
||||||
cursor.close()
|
cursor.close()
|
||||||
elif backend == "postgres":
|
elif backend == "postgres":
|
||||||
@ -139,31 +143,28 @@ async def init_engine(
|
|||||||
|
|
||||||
_session_factory = async_sessionmaker(_engine, expire_on_commit=False)
|
_session_factory = async_sessionmaker(_engine, expire_on_commit=False)
|
||||||
|
|
||||||
# Auto-create tables (dev convenience). Production should use Alembic.
|
# Schema bootstrap (hybrid):
|
||||||
from deerflow.persistence.base import Base
|
# - empty DB -> create_all + alembic stamp head
|
||||||
|
# - legacy DB -> create_all (baseline tables only, backfill) + alembic stamp baseline + upgrade head
|
||||||
# Import all models so Base.metadata discovers them.
|
# - already managed -> alembic upgrade head
|
||||||
# When no models exist yet (scaffolding phase), this is a no-op.
|
# Concurrency: Postgres advisory lock (true cross-process); SQLite uses an
|
||||||
try:
|
# in-process asyncio.Lock plus a 30s PRAGMA busy_timeout (also set on
|
||||||
import deerflow.persistence.models # noqa: F401
|
# alembic's own connections in env.py) -- multi-process SQLite bootstrap
|
||||||
except ImportError:
|
# is best-effort, gated by SQLite's natural file-level write lock.
|
||||||
# Models package not yet available — tables won't be auto-created.
|
# See deerflow.persistence.bootstrap for the full state machine.
|
||||||
# This is expected during initial scaffolding or minimal installs.
|
from deerflow.persistence.bootstrap import bootstrap_schema
|
||||||
logger.debug("deerflow.persistence.models not found; skipping auto-create tables")
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
async with _engine.begin() as conn:
|
await bootstrap_schema(_engine, backend=backend)
|
||||||
await conn.run_sync(Base.metadata.create_all)
|
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
if backend == "postgres" and "does not exist" in str(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)
|
await _auto_create_postgres_db(url)
|
||||||
# Rebuild engine against the now-existing database
|
# Rebuild engine against the now-existing database
|
||||||
await _engine.dispose()
|
await _engine.dispose()
|
||||||
_engine = create_async_engine(url, echo=echo, pool_size=pool_size, pool_pre_ping=True, json_serializer=_json_serializer)
|
_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)
|
_session_factory = async_sessionmaker(_engine, expire_on_commit=False)
|
||||||
async with _engine.begin() as conn:
|
await bootstrap_schema(_engine, backend=backend)
|
||||||
await conn.run_sync(Base.metadata.create_all)
|
|
||||||
else:
|
else:
|
||||||
raise
|
raise
|
||||||
|
|
||||||
|
|||||||
@ -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
|
||||||
@ -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)
|
||||||
@ -1,8 +1,15 @@
|
|||||||
"""Alembic environment for DeerFlow application tables.
|
"""Alembic environment for DeerFlow application tables.
|
||||||
|
|
||||||
ONLY manages DeerFlow's tables (runs, threads_meta, cron_jobs, users).
|
ONLY manages DeerFlow's tables (runs, threads_meta, feedback, users,
|
||||||
LangGraph's checkpointer tables are managed by LangGraph itself -- they
|
run_events, channel_connections, channel_credentials, channel_oauth_states,
|
||||||
have their own schema lifecycle and must not be touched by Alembic.
|
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
|
from __future__ import annotations
|
||||||
@ -15,6 +22,14 @@ from alembic import context
|
|||||||
from sqlalchemy.ext.asyncio import create_async_engine
|
from sqlalchemy.ext.asyncio import create_async_engine
|
||||||
|
|
||||||
from deerflow.persistence.base import Base
|
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.
|
# Import all models so metadata is populated.
|
||||||
try:
|
try:
|
||||||
@ -39,6 +54,7 @@ def run_migrations_offline() -> None:
|
|||||||
target_metadata=target_metadata,
|
target_metadata=target_metadata,
|
||||||
literal_binds=True,
|
literal_binds=True,
|
||||||
render_as_batch=True,
|
render_as_batch=True,
|
||||||
|
include_object=include_object,
|
||||||
)
|
)
|
||||||
with context.begin_transaction():
|
with context.begin_transaction():
|
||||||
context.run_migrations()
|
context.run_migrations()
|
||||||
@ -49,6 +65,7 @@ def do_run_migrations(connection):
|
|||||||
connection=connection,
|
connection=connection,
|
||||||
target_metadata=target_metadata,
|
target_metadata=target_metadata,
|
||||||
render_as_batch=True, # Required for SQLite ALTER TABLE support
|
render_as_batch=True, # Required for SQLite ALTER TABLE support
|
||||||
|
include_object=include_object,
|
||||||
)
|
)
|
||||||
with context.begin_transaction():
|
with context.begin_transaction():
|
||||||
context.run_migrations()
|
context.run_migrations()
|
||||||
@ -56,6 +73,25 @@ def do_run_migrations(connection):
|
|||||||
|
|
||||||
async def run_migrations_online() -> None:
|
async def run_migrations_online() -> None:
|
||||||
connectable = create_async_engine(config.get_main_option("sqlalchemy.url"))
|
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:
|
async with connectable.connect() as connection:
|
||||||
await connection.run_sync(do_run_migrations)
|
await connection.run_sync(do_run_migrations)
|
||||||
await connectable.dispose()
|
await connectable.dispose()
|
||||||
|
|||||||
@ -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"}
|
||||||
@ -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 ###
|
||||||
@ -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")
|
||||||
@ -4,7 +4,7 @@ from __future__ import annotations
|
|||||||
|
|
||||||
from datetime import UTC, datetime
|
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 sqlalchemy.orm import Mapped, mapped_column
|
||||||
|
|
||||||
from deerflow.persistence.base import Base
|
from deerflow.persistence.base import Base
|
||||||
@ -39,8 +39,7 @@ class RunRow(Base):
|
|||||||
lead_agent_tokens: Mapped[int] = mapped_column(default=0)
|
lead_agent_tokens: Mapped[int] = mapped_column(default=0)
|
||||||
subagent_tokens: Mapped[int] = mapped_column(default=0)
|
subagent_tokens: Mapped[int] = mapped_column(default=0)
|
||||||
middleware_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, server_default=text("'{}'"))
|
||||||
token_usage_by_model: Mapped[dict] = mapped_column(JSON, default=dict)
|
|
||||||
|
|
||||||
# Follow-up association
|
# Follow-up association
|
||||||
follow_up_to_run_id: Mapped[str | None] = mapped_column(String(64))
|
follow_up_to_run_id: Mapped[str | None] = mapped_column(String(64))
|
||||||
|
|||||||
78
backend/scripts/_autogen_revision.py
Normal file
78
backend/scripts/_autogen_revision.py
Normal file
@ -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()
|
||||||
75
backend/tests/blocking_io/test_persistence_bootstrap.py
Normal file
75
backend/tests/blocking_io/test_persistence_bootstrap.py
Normal file
@ -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()
|
||||||
@ -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
|
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
|
filesystem operation under test is the directory creation; if it regresses to run
|
||||||
directly on the event loop, Blockbuster raises `BlockingError` and this fails.
|
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
|
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.begin.return_value = begin_ctx
|
||||||
mock_engine.dispose = AsyncMock()
|
mock_engine.dispose = AsyncMock()
|
||||||
|
|
||||||
|
async def _noop_bootstrap(*_args, **_kwargs):
|
||||||
|
return None
|
||||||
|
|
||||||
with (
|
with (
|
||||||
patch.object(engine_mod, "create_async_engine", return_value=mock_engine),
|
patch.object(engine_mod, "create_async_engine", return_value=mock_engine),
|
||||||
patch.object(engine_mod, "async_sessionmaker", return_value=MagicMock()),
|
patch.object(engine_mod, "async_sessionmaker", return_value=MagicMock()),
|
||||||
patch("sqlalchemy.event.listens_for", _noop_listens_for),
|
patch("sqlalchemy.event.listens_for", _noop_listens_for),
|
||||||
|
patch(
|
||||||
|
"deerflow.persistence.bootstrap.bootstrap_schema",
|
||||||
|
new=_noop_bootstrap,
|
||||||
|
),
|
||||||
):
|
):
|
||||||
await engine_mod.init_engine(
|
await engine_mod.init_engine(
|
||||||
backend="sqlite",
|
backend="sqlite",
|
||||||
|
|||||||
@ -7,7 +7,7 @@ must stay identical.
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from sqlalchemy import Integer, String
|
from sqlalchemy import Integer, MetaData, String
|
||||||
from sqlalchemy.orm import Mapped, mapped_column
|
from sqlalchemy.orm import Mapped, mapped_column
|
||||||
|
|
||||||
from deerflow.persistence.base import Base, _column_keys
|
from deerflow.persistence.base import Base, _column_keys
|
||||||
@ -15,6 +15,10 @@ from deerflow.persistence.base import Base, _column_keys
|
|||||||
|
|
||||||
class _Widget(Base):
|
class _Widget(Base):
|
||||||
__tablename__ = "_widget_to_dict_test"
|
__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)
|
id: Mapped[int] = mapped_column(Integer, primary_key=True)
|
||||||
name: Mapped[str] = mapped_column(String(32))
|
name: Mapped[str] = mapped_column(String(32))
|
||||||
|
|||||||
102
backend/tests/test_persistence_autogen_script.py
Normal file
102
backend/tests/test_persistence_autogen_script.py
Normal file
@ -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)
|
||||||
636
backend/tests/test_persistence_bootstrap.py
Normal file
636
backend/tests/test_persistence_bootstrap.py
Normal file
@ -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}"
|
||||||
150
backend/tests/test_persistence_bootstrap_concurrency.py
Normal file
150
backend/tests/test_persistence_bootstrap_concurrency.py
Normal file
@ -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()
|
||||||
98
backend/tests/test_persistence_bootstrap_pg_lock.py
Normal file
98
backend/tests/test_persistence_bootstrap_pg_lock.py
Normal file
@ -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}"
|
||||||
121
backend/tests/test_persistence_bootstrap_regression.py
Normal file
121
backend/tests/test_persistence_bootstrap_regression.py
Normal file
@ -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()
|
||||||
116
backend/tests/test_persistence_bootstrap_sqlite_lock.py
Normal file
116
backend/tests/test_persistence_bootstrap_sqlite_lock.py
Normal file
@ -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
|
||||||
71
backend/tests/test_persistence_bootstrap_url.py
Normal file
71
backend/tests/test_persistence_bootstrap_url.py
Normal file
@ -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"
|
||||||
91
backend/tests/test_persistence_migrations_env.py
Normal file
91
backend/tests/test_persistence_migrations_env.py
Normal file
@ -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"
|
||||||
Loading…
x
Reference in New Issue
Block a user