mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-17 01:56:18 +00:00
* 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
655 lines
31 KiB
Python
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)
|