diff --git a/backend/packages/harness/deerflow/persistence/bootstrap.py b/backend/packages/harness/deerflow/persistence/bootstrap.py index c4f530ae7..92c9caad5 100644 --- a/backend/packages/harness/deerflow/persistence/bootstrap.py +++ b/backend/packages/harness/deerflow/persistence/bootstrap.py @@ -139,6 +139,45 @@ _BASELINE_TABLE_NAMES: frozenset[str] = frozenset( } ) +# ``test_baseline_index_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_INDEX_NAMES: frozenset[str] = frozenset( + { + # channel_connections + "idx_channel_connections_event_lookup", + "ix_channel_connections_owner_user_id", + "ix_channel_connections_provider", + "uq_channel_connection_active_identity", + # channel_conversations + "ix_channel_conversations_connection_id", + "ix_channel_conversations_owner_user_id", + "ix_channel_conversations_provider", + "ix_channel_conversations_thread_id", + # channel_oauth_states + "ix_channel_oauth_states_owner_user_id", + "ix_channel_oauth_states_provider", + # feedback + "ix_feedback_run_id", + "ix_feedback_thread_id", + "ix_feedback_user_id", + # run_events + "ix_events_run", + "ix_events_thread_cat_seq", + "ix_run_events_user_id", + # runs + "ix_runs_thread_id", + "ix_runs_thread_status", + "ix_runs_user_id", + # threads_meta + "ix_threads_meta_assistant_id", + "ix_threads_meta_user_id", + # users + "idx_users_oauth_identity", + "ix_users_email", + } +) + # 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 @@ -299,6 +338,41 @@ def _run_baseline_create_all_sync(sync_conn: Any) -> None: 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) + # ``create_all`` with ``checkfirst=True`` skips a table and all its + # subordinate ``Index`` objects when the table already exists. An index + # that was added to the ORM model after the table was first provisioned + # would therefore never be created, and because the legacy branch stamps + # ``0001_baseline`` before running upgrade, alembic's own + # ``batch_op.create_index`` for baseline-era indexes is skipped too. + # Explicitly creating every baseline-era ``Index`` on every baseline table + # (each with its own ``checkfirst=True``) guarantees each index exists + # regardless of whether its parent table was just created or already + # present. + # + # **Scope**: Only indexes in ``_BASELINE_INDEX_NAMES`` are created. + # ``table.indexes`` is the *current* ORM model's full index set, which + # includes post-baseline indexes added by later revisions (e.g. + # ``uq_runs_thread_active`` from 0004). Creating those prematurely would + # collide with their owning revision's data prerequisites (dedup steps, + # column migrations) and raise ``IntegrityError`` on legacy DBs. + # + # Post-baseline revisions that add an index to a baseline table must use + # the existing ``sa.inspect(bind).get_indexes(...)`` + ``if name not in + # existing`` guard pattern (see 0004_run_ownership.py:99-103), or a future + # ``safe_create_index`` helper -- mirroring ``safe_add_column``. + for table in baseline_tables: + for idx in table.indexes: + if idx.name not in _BASELINE_INDEX_NAMES: + continue + try: + idx.create(sync_conn, checkfirst=True) + except Exception: + logger.warning( + "bootstrap: failed to create baseline index %r on %r -- the DB may contain rows that violate the index constraint. Address the duplicate data, then re-run bootstrap.", + idx.name, + table.name, + ) + def _stamp(cfg: AlembicConfig, revision: str) -> None: """Synchronous alembic stamp; callers must wrap in ``asyncio.to_thread``.""" diff --git a/backend/tests/test_persistence_bootstrap.py b/backend/tests/test_persistence_bootstrap.py index 0a9d8b0f6..257208c3b 100644 --- a/backend/tests/test_persistence_bootstrap.py +++ b/backend/tests/test_persistence_bootstrap.py @@ -32,6 +32,7 @@ from sqlalchemy.ext.asyncio import create_async_engine 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, @@ -558,6 +559,35 @@ async def test_baseline_table_names_constant_matches_0001(tmp_path: Path) -> Non 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 @@ -591,6 +621,197 @@ async def test_legacy_backfill_skips_non_baseline_tables(tmp_path: Path) -> None 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) # ---------------------------------------------------------------------------