mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-21 10:15:47 +00:00
VanzerenandGitHub
1769b2de0d
* fix: read run stop_reason from runtime context * fix: address review feedback for #4188 stop_reason integration - migration 0005: use safe_add_column for consistency and drift detection - worker: clear runtime.context stop_reason at start of each _stream_once turn so a clean continuation doesn't inherit a prior cap reason - tests: replace circular unit test with real middleware integration tests that exercise LoopDetectionMiddleware._apply and TokenBudgetMiddleware._apply through the worker, proving the full middleware → runtime.context → persist pipeline * fix(test): resume conftest * fix: stamp stop_reason in all guard middlewares, fix clearing semantics
869 lines
40 KiB
Python
869 lines
40 KiB
Python
"""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_INDEX_NAMES,
|
||
_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 = "0005_run_stop_reason"
|
||
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 _runs_index_names(engine) -> set[str]:
|
||
async with engine.connect() as conn:
|
||
return await conn.run_sync(lambda c: {ix["name"] for ix in sa.inspect(c).get_indexes("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
|
||
# The partial unique index on (thread_id WHERE status IN pending/running)
|
||
# must exist on a fresh DB because the empty-branch stamps head without
|
||
# running migrations, so the index has to come from ``Base.metadata``.
|
||
indexes = await _runs_index_names(engine)
|
||
assert "uq_runs_thread_active" in indexes, indexes
|
||
assert "ix_runs_lease" in indexes, indexes
|
||
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_baseline_index_names_constant_matches_0001(tmp_path: Path) -> None:
|
||
"""Guard: ``_BASELINE_INDEX_NAMES`` must match the set of indexes that
|
||
``0001_baseline.upgrade()`` actually creates. Editing 0001 (or adding an
|
||
index to the ORM model without a corresponding migration guard) without
|
||
updating this constant fires this test.
|
||
"""
|
||
engine = create_async_engine(_url(tmp_path))
|
||
try:
|
||
cfg = _get_alembic_config(engine)
|
||
await asyncio.to_thread(_upgrade, cfg, "0001_baseline")
|
||
|
||
async with engine.connect() as conn:
|
||
all_indexes: set[str] = set()
|
||
for table_name in _BASELINE_TABLE_NAMES:
|
||
indexes = await conn.run_sync(lambda c, t=table_name: {ix["name"] for ix in sa.inspect(c).get_indexes(t)})
|
||
all_indexes.update(indexes)
|
||
|
||
# alembic_version also gets an index from alembic's own table --
|
||
# filter it out like the table guard filters alembic_version.
|
||
all_indexes.discard("ix_alembic_version")
|
||
|
||
assert all_indexes == _BASELINE_INDEX_NAMES, (
|
||
f"_BASELINE_INDEX_NAMES drifted from 0001_baseline.upgrade()'s output:\nonly-in-0001={sorted(all_indexes - _BASELINE_INDEX_NAMES)}\nonly-in-constant={sorted(_BASELINE_INDEX_NAMES - all_indexes)}"
|
||
)
|
||
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)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Legacy backfill: an index added to the ORM model after the table was first
|
||
# provisioned must also be created by the backfill helper, not only when the
|
||
# table itself is missing. ``create_all`` with ``checkfirst=True`` skips a
|
||
# table and all its ``Index`` objects when the table already exists (it only
|
||
# checks the table, not each index). The backfill must therefore create every
|
||
# baseline table's indexes explicitly, covering the case where the table is
|
||
# present but one of its indexes was added in a later model change.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
async def _channel_connections_index_names(engine) -> set[str]:
|
||
async with engine.connect() as conn:
|
||
return await conn.run_sync(lambda c: {ix["name"] for ix in sa.inspect(c).get_indexes("channel_connections")})
|
||
|
||
|
||
@asyncio_test
|
||
async def test_legacy_backfill_creates_missing_index_on_existing_table(tmp_path: Path) -> None:
|
||
engine = create_async_engine(_url(tmp_path))
|
||
try:
|
||
# Simulate a legacy DB where ``channel_connections`` was provisioned
|
||
# before the ``uq_channel_connection_active_identity`` partial unique
|
||
# index was added to the model -- create everything, then drop just
|
||
# that index.
|
||
async with engine.begin() as conn:
|
||
await conn.run_sync(Base.metadata.create_all)
|
||
async with engine.begin() as conn:
|
||
await conn.execute(sa.text("DROP INDEX IF EXISTS uq_channel_connection_active_identity"))
|
||
|
||
# Confirm the index is gone.
|
||
indexes_before = await _channel_connections_index_names(engine)
|
||
assert "uq_channel_connection_active_identity" not in indexes_before, indexes_before
|
||
|
||
# Run the backfill helper -- this is what the legacy branch calls.
|
||
async with engine.begin() as conn:
|
||
await conn.run_sync(_run_baseline_create_all_sync)
|
||
|
||
# The index must now exist.
|
||
indexes_after = await _channel_connections_index_names(engine)
|
||
assert "uq_channel_connection_active_identity" in indexes_after, indexes_after
|
||
finally:
|
||
await engine.dispose()
|
||
|
||
|
||
@asyncio_test
|
||
async def test_legacy_backfill_idempotent_when_index_already_exists(tmp_path: Path) -> None:
|
||
"""The explicit index-creation pass must not raise when the index is already present."""
|
||
engine = create_async_engine(_url(tmp_path))
|
||
try:
|
||
async with engine.begin() as conn:
|
||
await conn.run_sync(Base.metadata.create_all)
|
||
|
||
indexes_before = await _channel_connections_index_names(engine)
|
||
assert "uq_channel_connection_active_identity" in indexes_before
|
||
|
||
# Running the backfill helper over an already-correct DB must succeed.
|
||
async with engine.begin() as conn:
|
||
await conn.run_sync(_run_baseline_create_all_sync)
|
||
|
||
indexes_after = await _channel_connections_index_names(engine)
|
||
assert "uq_channel_connection_active_identity" in indexes_after
|
||
finally:
|
||
await engine.dispose()
|
||
|
||
|
||
@asyncio_test
|
||
async def test_legacy_backfill_rejects_post_baseline_indexes(tmp_path: Path) -> None:
|
||
"""Regression: the index loop must not create post-baseline indexes (e.g.
|
||
``uq_runs_thread_active`` from 0004) that have data prerequisites. If the
|
||
backfill creates a partial unique index before the owning revision runs its
|
||
dedup step, duplicate rows in the legacy DB raise ``IntegrityError`` and
|
||
crash bootstrap.
|
||
|
||
We seed two active (status='pending') runs on the same thread -- the exact
|
||
case revision 0004's ``_dedupe_active_runs_per_thread`` was written for --
|
||
and assert the backfill helper succeeds without error, then confirm
|
||
0004 still runs correctly (dedup + index creation) during upgrade.
|
||
"""
|
||
engine = create_async_engine(_url(tmp_path))
|
||
try:
|
||
# 1. Simulate a legacy DB: create baseline schema (0001 only), then
|
||
# seed violating data that a post-baseline partial unique index
|
||
# would reject.
|
||
cfg = _get_alembic_config(engine)
|
||
await asyncio.to_thread(_upgrade, cfg, "0001_baseline")
|
||
|
||
async with engine.begin() as conn:
|
||
# Insert two pending runs for the same thread -- violates the
|
||
# partial UNIQUE constraint uq_runs_thread_active that 0004 adds.
|
||
for i in range(2):
|
||
await conn.execute(
|
||
sa.text(
|
||
"INSERT INTO runs "
|
||
"(run_id, thread_id, user_id, assistant_id, status, "
|
||
"multitask_strategy, metadata_json, kwargs_json, "
|
||
"message_count, total_input_tokens, total_output_tokens, "
|
||
"total_tokens, llm_call_count, lead_agent_tokens, "
|
||
"subagent_tokens, middleware_tokens, token_usage_by_model, "
|
||
"created_at, updated_at) "
|
||
"VALUES (:run_id, 'thread-1', 'u1', 'asst-1', 'pending', "
|
||
"'reject', '{}', '{}', 0, 0, 0, 0, 0, 0, 0, 0, '{}', "
|
||
":created_at, :updated_at)"
|
||
),
|
||
{
|
||
"run_id": f"run-{i}",
|
||
"created_at": f"2026-01-01T00:00:0{i}",
|
||
"updated_at": f"2026-01-01T00:00:0{i}",
|
||
},
|
||
)
|
||
|
||
# 2. Run the backfill helper -- must NOT crash on duplicate data.
|
||
async with engine.begin() as conn:
|
||
await conn.run_sync(_run_baseline_create_all_sync)
|
||
|
||
# 3. Stamp baseline and upgrade to head -- 0004 should dedupe then
|
||
# create the index cleanly.
|
||
await asyncio.to_thread(_upgrade, cfg, "0001_baseline") # stamp-like: 0001 already ran
|
||
await asyncio.to_thread(_upgrade, cfg, "head")
|
||
|
||
# 4. After upgrade, the index must exist and only one active run
|
||
# should remain (0004's dedup cancelled the other).
|
||
async with engine.connect() as conn:
|
||
indexes = await conn.run_sync(lambda c: {ix["name"] for ix in sa.inspect(c).get_indexes("runs")})
|
||
assert "uq_runs_thread_active" in indexes, "0004 should have created uq_runs_thread_active after dedup"
|
||
|
||
# 0004's dedup cancels all-but-one active run per thread.
|
||
remaining = (await conn.execute(sa.text("SELECT COUNT(*) FROM runs WHERE thread_id='thread-1' AND status='pending'"))).scalar()
|
||
assert remaining == 1, f"Expected 1 pending run after 0004 dedup, got {remaining}"
|
||
finally:
|
||
await engine.dispose()
|
||
|
||
|
||
@asyncio_test
|
||
async def test_legacy_backfill_duplicate_channel_connections_does_not_crash(
|
||
tmp_path: Path,
|
||
) -> None:
|
||
"""The target index ``uq_channel_connection_active_identity`` is a baseline
|
||
partial UNIQUE. Seeding duplicate non-revoked channel connections must
|
||
NOT crash the backfill with ``IntegrityError``, since no revision dedupes
|
||
channel connections (unlike runs). A crash would brick bootstrap with no
|
||
self-heal path.
|
||
|
||
The backfill catches the creation error and logs a warning; the operator
|
||
must fix the duplicate data and re-run.
|
||
"""
|
||
engine = create_async_engine(_url(tmp_path))
|
||
try:
|
||
cfg = _get_alembic_config(engine)
|
||
await asyncio.to_thread(_upgrade, cfg, "0001_baseline")
|
||
|
||
# Drop the index so the backfill has work to do.
|
||
async with engine.begin() as conn:
|
||
await conn.execute(sa.text("DROP INDEX IF EXISTS uq_channel_connection_active_identity"))
|
||
|
||
# Seed duplicate non-revoked connections with the same identity.
|
||
async with engine.begin() as conn:
|
||
for i in range(2):
|
||
await conn.execute(
|
||
sa.text(
|
||
"INSERT INTO channel_connections "
|
||
"(id, owner_user_id, provider, status, external_account_id, "
|
||
"workspace_id, scopes_json, capabilities_json, metadata_json, "
|
||
"created_at, updated_at) "
|
||
"VALUES (:cid, :owner, 'slack', 'active', 'ext-1', 'ws-1', "
|
||
"'[]', '{}', '{}', :created_at, :updated_at)"
|
||
),
|
||
{
|
||
"cid": f"cc-{i}",
|
||
# The baseline owner+identity unique constraint requires
|
||
# distinct owners; the partial active-identity index does not.
|
||
"owner": f"u{i}",
|
||
"created_at": "2026-01-01T00:00:00",
|
||
"updated_at": "2026-01-01T00:00:00",
|
||
},
|
||
)
|
||
|
||
# Backfill must not crash even with duplicate data.
|
||
async with engine.begin() as conn:
|
||
await conn.run_sync(_run_baseline_create_all_sync)
|
||
|
||
# The index won't exist (duplicate data prevents creation), but
|
||
# bootstrap must have survived -- the operator can fix data and retry.
|
||
async with engine.connect() as conn:
|
||
indexes = await conn.run_sync(lambda c: {ix["name"] for ix in sa.inspect(c).get_indexes("channel_connections")})
|
||
# Index missing is expected when data violates the constraint.
|
||
# The backfill logged a warning instead of crashing.
|
||
assert "uq_channel_connection_active_identity" not in indexes
|
||
finally:
|
||
await engine.dispose()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# _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}"
|