mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-16 17:46:20 +00:00
06c827903a
* feat(persistence): expand thread incarnation storage Add nullable thread and MCP task incarnation columns while preserving mixed-version writes. New thread records receive stable incarnation IDs, and new task rows copy the matching owned or shared thread incarnation without changing any read, claim, session, or deletion behavior. * test(persistence): pin incarnation rollback compatibility * test(api): pin internal thread response boundary * fix(persistence): rebase incarnation rollout after projects --------- Co-authored-by: CorgiBoyG <CorgiBoyG@users.noreply.github.com>
701 lines
35 KiB
Python
701 lines
35 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 compatibility window, a rollback-floor SQLite process also
|
|
re-reads ``alembic_version`` after an Alembic ``CommandError``. It recovers
|
|
only when another process advanced the file to the explicitly reviewed
|
|
incarnation revision and that revision is absent from the rollback binary's
|
|
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 the 0020 rollback-floor 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. Tests cross-pin
|
|
# this revision id and schema shape; amending that DDL requires re-auditing old
|
|
# repository reads and writes before this exception remains valid. The
|
|
# exception also requires every rollback-floor ORM table and column.
|
|
#
|
|
# The revision id was deployed briefly as an out-of-tree child of 0018 before
|
|
# becoming the in-tree successor to ``0020_threads_meta_project_id``. Current
|
|
# and future binaries that know this reused stamp validate the fixed
|
|
# canonical-0019 floor before migrating onward. This code also uses that fixed
|
|
# floor when tests remove 0019 from the local revision set to exercise the
|
|
# unknown-revision path. The published 0020 rollback binary instead validates
|
|
# its own ORM floor before skipping the revision.
|
|
_FORWARD_COMPATIBLE_REVISION = "0019_thread_incarnations"
|
|
|
|
# Fixed table/column floor produced by the canonical in-tree 0019 revision.
|
|
# Never derive this from Base.metadata: future ORM fields belong to later
|
|
# migrations and must not be required before those migrations run.
|
|
_CANONICAL_0019_SCHEMA_FLOOR: dict[str, frozenset[str]] = {
|
|
"agents": frozenset("config created_at id name soul updated_at user_id".split()),
|
|
"channel_connections": frozenset(
|
|
"bot_user_id capabilities_json created_at external_account_id external_account_name id last_error_at last_seen_at metadata_json owner_user_id provider scopes_json status updated_at workspace_id workspace_name".split()
|
|
),
|
|
"channel_conversations": frozenset("connection_id created_at external_conversation_id external_topic_id id owner_user_id provider thread_id updated_at".split()),
|
|
"channel_credentials": frozenset("connection_id encrypted_access_token encrypted_extra_json encrypted_refresh_token expires_at refresh_expires_at token_type updated_at version".split()),
|
|
"channel_oauth_states": frozenset("code_verifier_encrypted consumed_at created_at expires_at metadata_json nonce_hash owner_user_id provider redirect_after requested_scopes_json state_hash".split()),
|
|
"feedback": frozenset("comment created_at feedback_id message_id rating run_id thread_id user_id".split()),
|
|
"managed_subagents": frozenset("created_at definition id name updated_at".split()),
|
|
"mcp_tasks": frozenset(
|
|
"""cancel_attempt_count cancel_requested_at completed_at consecutive_poll_error_count created_at dispatch_attempt dispatch_event dispatch_version driver_data driver_name error event_fingerprint event_version id input_required
|
|
last_cancel_error last_poll_error last_polled_at lease_expires_at lease_owner next_cancel_at next_notification_at next_poll_at notification_attempt_count notification_error notification_lease_expires_at
|
|
notification_lease_owner notification_run_id notification_status notified_version poll_attempt_count remote_task_id result result_artifact result_preview result_truncated run_id server_name status task_name
|
|
thread_id thread_incarnation tool_call_id updated_at user_id""".split()
|
|
),
|
|
"personal_access_tokens": frozenset("created_at expires_at id last_used_at name revoked_at scopes token_digest user_id".split()),
|
|
"projects": frozenset("created_at id instructions name presentation status updated_at user_id".split()),
|
|
"run_events": frozenset("category content created_at event_metadata event_type id run_id seq thread_id user_id".split()),
|
|
"runs": frozenset(
|
|
"""assistant_id cancel_action cancel_requested_at created_at error first_human_message follow_up_to_run_id idempotency_key kwargs_json last_ai_message lead_agent_tokens lease_expires_at llm_call_count
|
|
message_count metadata_json middleware_tokens model_name multitask_strategy operation_kind owner_worker_id run_id status stop_reason subagent_tokens thread_id token_usage_by_model total_input_tokens
|
|
total_output_tokens total_tokens updated_at user_id""".split()
|
|
),
|
|
"scheduled_task_runs": frozenset("attempt_count created_at error finished_at id lease_expires_at lease_owner run_id scheduled_for started_at status task_id thread_id trigger".split()),
|
|
"scheduled_tasks": frozenset(
|
|
"""assistant_id context_mode created_at id last_error last_run_at last_run_id last_thread_id lease_expires_at lease_owner next_run_at overlap_policy prompt run_count schedule_spec schedule_type
|
|
status thread_id timezone title updated_at user_id""".split()
|
|
),
|
|
"subagent_batch_items": frozenset(
|
|
"""acceptance_criteria acceptance_verdict attempt batch_id cancel_requested_at completed_at created_at error id item_key lease_expires_at lease_owner model_name position prompt result result_preview
|
|
result_truncated started_at status stop_reason token_usage updated_at""".split()
|
|
),
|
|
"subagent_batches": frozenset("completed_at created_at execution_spec id max_attempts max_live_items max_running_items run_id status subagent_type submission_key thread_id title tool_call_id total_items updated_at user_id".split()),
|
|
"threads_meta": frozenset("assistant_id created_at display_name incarnation metadata_json project_id status thread_id updated_at user_id".split()),
|
|
"users": frozenset("created_at email id needs_setup oauth_id oauth_provider password_hash system_role token_version".split()),
|
|
"webhook_deliveries": frozenset("channel chat_id first_seen message_id workspace_id".split()),
|
|
}
|
|
|
|
# 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 fixed canonical-0019 schema before acting on its stamp.
|
|
|
|
This is a presence check, not a general schema compatibility proof. The
|
|
allowlisted additive DDL still needs its separate read/write audit. The
|
|
floor remains independent of current ORM metadata so future revisions can
|
|
add their own schema after this validation succeeds.
|
|
"""
|
|
inspector = sa_inspect(sync_conn)
|
|
tables = set(inspector.get_table_names())
|
|
missing = []
|
|
for name, required_columns in sorted(_CANONICAL_0019_SCHEMA_FLOOR.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_name in sorted(required_columns - 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":
|
|
# The same revision id once named a different out-of-tree schema.
|
|
# Validate canonical 0019 before upgrading a known revision or
|
|
# accepting it through the unknown-revision path that current
|
|
# tests use to simulate old 0020.
|
|
if database_revision == _FORWARD_COMPATIBLE_REVISION:
|
|
async with engine.connect() as conn:
|
|
await conn.run_sync(_validate_forward_schema)
|
|
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:
|
|
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)
|