Zeren Wang 5951c89b5b
feat(projects): project workspaces with scoped chats and thread membership (#5265)
* feat(projects): project workspaces with scoped chats and thread membership

Backend:
- projects table model and migration; fail-closed ProjectRepository with
  ownership checks, CRUD/archive/restore/delete router, and atomic thread
  move between projects
- threads_meta.project_id column exposed as reserved deerflow_project_id
  metadata; project-aware thread create/search with pagination bounds and
  membership echoed in create responses
- first-run admission assigns the project only at genuine first run, seeded
  at write time and dropped when invalid; serialized against project
  deletion and thread assignment
- branch creation inherits the source thread's project membership (an
  archived/deleted project degrades the branch to unassigned instead of
  failing the request)

Frontend:
- projects data layer, thread move API, and sidebar projects section with
  flat/grouped modes, archived-project threads, and stable virtual-list
  offsets
- project detail page with project-scoped new chat
  (/workspace/chats/new?project=) and paginated thread list
- move-to-project thread menu, new-project dialog, archived-project gates
- project-scoped new chats pre-create the thread with membership before the
  first submit or /goal set, so runs never proceed outside the project
- goal-set preparation is fenced against conversation switches: a stale
  continuation is dropped instead of saving the goal or launching the
  abandoned submission on the newly opened conversation
- project thread lists join thread lifecycle invalidations (stop, pin) so
  an open project page never keeps stale titles, recency, or pagination

* fix(chats): keep archive undo toast when the sidebar row unmounts

The archive success toast was fired from per-mutate callbacks passed to
mutation.mutate. React Query drops those handlers when the observer
component unmounts before the mutation settles; archiving the open chat
removes its sidebar row mid-flight, so the undo toast never appeared and
the e2e archive-undo test timed out waiting for it.

Move the success/error handlers to the mutation level (useArchiveThread
options, same pattern as useMoveThreadToProject) where callbacks are
delivered even after the originating row unmounts.

* fix(projects): pin project thread listing contract and exclude archived chats

GET /api/projects/{id}/threads returned the thread store row verbatim
(list[dict], no response_model): user_id/assistant_id leaked, any future
ThreadMetaRow column would auto-leak, and the OpenAPI schema was empty.
Return a narrow ProjectThreadResponse (the exact fields ProjectThread
declares) with the same metadata secret redaction the surrounding thread
endpoints get from _MetadataRedactingResponse.

The listing also ran search() without the archived filter, so a retired
chat rendered as a normal row on the project page while the sidebar hid
it. Search archived=False to mirror the sidebar's archived:false lists;
restore stays on the global Archived tab.

Both regressions pinned by new router tests: wire-shape allowlist and
archived-member exclusion.

* docs(migrations): record the 0019/0020 chain against the bootstrap reservation

The tree now chains 0018 -> 0019_projects -> 0020_threads_meta_project_id,
so migrations/AGENTS.md was stale twice over: the revision index stopped at
0018 and the rolling-forward section still claimed the tree 'deliberately
remains at 0018'.

Document the new head and record the intentional numeric-prefix reuse of
0019: 0019_projects is in-chain while 0019_thread_incarnations stays the
reserved, allowlisted out-of-tree rollout id. The owning rollout revision
must re-parent onto this tree's head when it merges so alembic never sees
two heads off 0018; bootstrap.py now cross-references that note next to
_FORWARD_COMPATIBLE_REVISION.

* fix(chats): invalidate project thread lists on archive/restore

useArchiveThread refreshed the infinite sidebar cache, threads/search and
the per-thread metadata cache but not the project-scoped list
([...PROJECTS_QUERY_KEY, 'threads', id]) this PR adds — the one thread
mutation not wired to that key, after usePinThread, useRenameThread,
useDeleteThread, useMoveThreadToProject and invalidateStoppedThreadCaches.

An archive from a sidebar row while a project page is open therefore left
the archived chat rendered as a normal row until remount (and undo left it
missing). Invalidate the prefix in the mutation-level success handler.

Regression test asserts the project-list prefix is invalidated on success.

* fix(projects): fetch project discovery only in grouped sidebar mode

RecentChatList mounted two useProjects queries per sidebar render, but
knownProjectIds is consumed only by the grouped-mode exclusion filter; in
the default flat mode every page load paid two GET /api/projects?status=
round trips for data nothing read. Gate both queries on grouped mode —
GroupedProjectList fetches the same keys when the toggle is on and
TanStack dedupes the observers.

Also set retry: false on useProject: a deleted or foreign project 404s
deterministically, and the page renders a dedicated not-found state for
it, so the default 1s/2s/4s retry backoff kept deep links in 'loading'
for ~7s before that state appeared. Matches useThreadMetadata /
useThreadTokenUsage.

* fix(threads): fail closed on project-scoped create in memory mode

MemoryThreadMetaStore.create accepted project_id and silently ignored it,
making memory mode the one membership path that fails open: POST
/api/threads with a project id returned 200 and the run started
unassigned, violating the invariant that a run never proceeds outside the
selected project (the SQL store raises ProjectNotAssignableError inside
the insert transaction for the same request).

Raise ProjectNotAssignableError whenever project_id is present so the
router's existing 404 mapping applies, the frontend keeps the composer
text for a retry, and memory mode behaves exactly like SQL mode.
set_project already reports rejection; create now matches it.

Store-level test (raises, nothing persisted, project filter stays empty,
unscoped creates still work) plus a router-level test asserting the 404
and that no row is left behind.

* fix(projects): window the project page thread list

ProjectThreadsSection rendered every loaded page as a plain Link row, so a
long-lived project accumulated unbounded DOM on the page's scroll surface:
each load-more appended another 100 rows and every formatTimeAgo tick
re-rendered the whole list.

Reuse VirtualThreadList (now generic over any row shape with a
thread_id), pointing its scroll parent at this page's ScrollArea viewport
via the shared [data-slot="scroll-area-viewport"] selector used by
/workspace/chats; under the 60-row threshold it falls back to the plain
render, so small projects are unchanged.

* fix(projects): restore row dividers and pin them with a render test

The row class template literal concatenated transition-colors directly
with the conditional border-b token, so non-final rows rendered the
invalid class 'transition-colorsborder-b' and lost both the divider and
the transition. Compose the row classes with cn() and a boolean guard
instead.

The section moved out of page.tsx into a testable component so the row
markup finally has coverage: a DOM test asserts every row except the
final data row carries border-b (index-based, not last: — correct under
virtualization where the last mounted row is not the last data row), and
the untitled fallback plus load-more button render for a partial page.

* fix(projects): validate forward schemas and fence membership reads
2026-09-08 17:00:26 +08:00

655 lines
31 KiB
Python

"""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 (one locally known revision) | ``alembic upgrade head`` |
| reviewed forward revision with local columns | warn and skip migration |
| unknown, empty, or multiple revision rows | refuse to start |
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.
During the 0018-to-0019 compatibility window, an old SQLite process also
re-reads ``alembic_version`` after an Alembic ``CommandError``. It recovers
only when another process advanced the file to the explicitly reviewed 0019
and 0019 is still absent from the local migration tree; every other migration
failure remains fatal.
``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 alembic.util.exc import CommandError
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
_KNOWN_REVISIONS: frozenset[str] | None = None
# One additive revision may be present when an older Gateway starts during the
# thread-incarnation rollout. This allowlist was reviewed only for revision
# ``0019_thread_incarnations`` adding nullable VARCHAR(32) columns
# ``threads_meta.incarnation`` and ``mcp_tasks.thread_incarnation`` without a
# server default, table, index, constraint, or data backfill. The owning 0019
# change must cross-pin this revision id and schema shape in tests. Amending
# that DDL requires re-auditing old-repository reads and writes before this
# exception remains valid. The exception also requires every current ORM table
# and column: the original 0018 + incarnation columns shape lacks projects and
# is no longer compatible with this build. Both skip paths validate that floor.
# Note: this tree's own chain already carries ``0019_projects`` /
# ``0020_threads_meta_project_id`` off ``0018_oauth_identity_pg_partial``; the
# ``0019_`` numeric prefix is intentionally reused. When the owning rollout
# revision merges it must re-parent onto the current head (see
# ``migrations/AGENTS.md``) so alembic never sees two heads off 0018.
_FORWARD_COMPATIBLE_REVISION = "0019_thread_incarnations"
# 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",
}
)
# ``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
# 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, *, postgres_schema: str = "") -> 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.
When *postgres_schema* is set it is forwarded as the ``deerflow_pg_schema``
main option so ``env.py`` can pin its alembic-spawned engine's
``search_path`` to the same schema the app engine uses. Without it,
alembic's own engine -- built from the bare URL -- would create
``alembic_version`` and all migration DDL in the default (``public``)
schema while the app tables land in the custom schema.
"""
cfg = AlembicConfig()
cfg.set_main_option("script_location", str(_MIGRATIONS_DIR))
cfg.set_main_option("sqlalchemy.url", _alembic_safe_url(engine))
if postgres_schema:
cfg.set_main_option("deerflow_pg_schema", postgres_schema)
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 _get_known_revisions() -> frozenset[str]:
"""Return every revision id available in the local migration tree."""
global _KNOWN_REVISIONS
if _KNOWN_REVISIONS is None:
cfg = AlembicConfig()
cfg.set_main_option("script_location", str(_MIGRATIONS_DIR))
script = ScriptDirectory.from_config(cfg)
_KNOWN_REVISIONS = frozenset(revision.revision for revision in script.walk_revisions())
return _KNOWN_REVISIONS
def _get_revision_metadata() -> tuple[str, frozenset[str]]:
"""Load the local head and revision set off the event loop."""
return _get_head_revision(), _get_known_revisions()
async def _read_database_revision(conn: Any) -> str:
"""Read and validate the database's single alembic revision row."""
result = await conn.execute(text("SELECT version_num FROM alembic_version"))
rows = list(result.scalars())
if len(rows) != 1:
raise RuntimeError(f"bootstrap: expected exactly one alembic_version row, found {len(rows)}")
revision = rows[0]
if not isinstance(revision, str) or not revision:
raise RuntimeError("bootstrap: alembic_version contains an empty revision")
return revision
def _validate_forward_schema(sync_conn: Any) -> None:
"""Require the local repository schema before skipping unknown migrations.
This is a presence check, not a general schema compatibility proof. The
allowlisted additive DDL still needs its separate read/write audit. Derive
the local floor from ORM metadata so a new mapped column cannot silently
invalidate the existing exception again.
"""
import deerflow.persistence.models # noqa: F401
from deerflow.persistence.base import Base
inspector = sa_inspect(sync_conn)
tables = set(inspector.get_table_names())
missing = []
for name, table in sorted(Base.metadata.tables.items()):
if name not in tables:
missing.append(name)
continue
columns = {column["name"] for column in inspector.get_columns(name)}
missing.extend(f"{name}.{column.name}" for column in table.columns if column.name not in columns)
if missing:
raise RuntimeError(
f"bootstrap: revision {_FORWARD_COMPATIBLE_REVISION!r} is missing required local schema: {', '.join(missing)}; refusing to start. See docs/database-forward-revision-recovery.md for the audited offline migration path."
)
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)
# ``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``."""
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, postgres_schema: 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``.
*postgres_schema*, when set, is forwarded to the alembic config so the
alembic-spawned engine pins its ``search_path`` to that schema. The target
schema must already exist (``init_engine`` issues ``CREATE SCHEMA`` before
calling this). Ignored for non-postgres backends.
"""
head, known_revisions = await asyncio.to_thread(_get_revision_metadata)
cfg = _get_alembic_config(engine, postgres_schema=postgres_schema if backend == "postgres" else "")
async with _bootstrap_lock(engine, backend=backend):
async with engine.connect() as conn:
state = await conn.run_sync(_reflect_state)
database_revision = await _read_database_revision(conn) if state["has_alembic_version"] else None
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":
if database_revision in known_revisions:
logger.info(
"bootstrap: branch=versioned revision=%s -> upgrade head (%s)",
database_revision,
head,
)
try:
await asyncio.to_thread(_upgrade, cfg, "head")
except CommandError:
# SQLite has no cross-process bootstrap mutex. Another
# process may advance 0018 to the reviewed 0019 after this
# process reads the version but before Alembic starts.
# Do not apply this recovery once 0019 belongs to the local
# tree: a new binary's migration failure must stay fatal.
if backend != "sqlite" or _FORWARD_COMPATIBLE_REVISION in known_revisions:
raise
async with engine.connect() as conn:
current_revision = await _read_database_revision(conn)
if current_revision != _FORWARD_COMPATIBLE_REVISION:
raise
await conn.run_sync(_validate_forward_schema)
logger.warning(
"bootstrap: database advanced concurrently to explicitly forward-compatible revision %s; skipping the stale local upgrade",
current_revision,
)
elif database_revision == _FORWARD_COMPATIBLE_REVISION:
async with engine.connect() as conn:
await conn.run_sync(_validate_forward_schema)
logger.warning(
"bootstrap: database revision %s is explicitly forward-compatible with local head %s and has its required tables and columns; skipping migration",
database_revision,
head,
)
else:
raise RuntimeError(f"bootstrap: database revision {database_revision!r} is not known to this build (local head {head!r}); refusing to start")
else: # pragma: no cover -- defensive
raise RuntimeError(f"bootstrap: unhandled decision {decision!r}")
logger.info("bootstrap: complete (backend=%s)", backend)