mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-01 10:56:02 +00:00
e3ddf4e331
* feat(persistence): wire alembic migrations + bootstrap schema on startup Closes #3682. Pre-#3658 DBs lack the `runs.token_usage_by_model` column because alembic was never wired up — startup only ran `create_all`, which never ALTERs existing tables. Adds a hybrid bootstrap in FastAPI lifespan (replaces bare `create_all`): - empty DB → create_all + stamp head - legacy DB → stamp 0001_baseline + upgrade head - versioned DB → upgrade head Concurrency: Postgres `pg_advisory_lock` (cross-process); SQLite per-engine `asyncio.Lock` + 30s `PRAGMA busy_timeout` on both prod and alembic engines. Column revisions use `safe_add_column` / `safe_drop_column` idempotent helpers as fallback. Other bits: - 0001 baseline (chain root) + 0002 add `runs.token_usage_by_model` - `include_object` filter so alembic ignores LangGraph checkpointer tables - `make migrate-rev MSG="..."` for authoring new revisions (no migrate/stamp targets — startup is the only execution path) - Tests: three-branch decision, concurrency, #3682 regression, env filter, blocking-IO gate anchor - CLAUDE.md: new "Schema migrations" section * fix(style): fix lint error * perf(persistence): address review feedback on alembic bootstrap Behavioural fixes - _SQLITE_LOCKS now keyed via WeakKeyDictionary so id-reuse after GC cannot return a stale, loop-bound lock and the cache cannot leak one entry per disposed engine. - safe_add_column compares nullable / server_default against the desired column when the name already exists and emits a warning on drift, surfacing manual-ALTER workarounds instead of silently no-op'ing. - _postgres_lock issues SET LOCAL idle_in_transaction_session_timeout=0 before pg_advisory_lock, so managed Postgres cannot kill the idle lock-holding session mid-upgrade and silently release the advisory lock. - legacy branch now backfills missing baseline tables via a restricted create_all (Base.metadata.create_all scoped to _BASELINE_TABLE_NAMES). Restores pre-#1930 upgraders whose channel_* tables were never provisioned, without pre-empting future create_table revisions for newly-added models. Schema parity - runs.token_usage_by_model gains server_default=text("'{}'") in both the ORM model and the 0001_baseline create_table, matching what 0002 adds via ALTER. create_all and alembic-upgrade paths now produce identical column definitions. - New parity test compares Base.metadata.create_all output against a pure alembic upgrade base->head, asserting column-set, nullable, and server_default agree across all tables (normalized through the same helper safe_add_column's drift check uses). Guards - test_baseline_table_names_constant_matches_0001 pins _BASELINE_TABLE_NAMES to 0001_baseline.upgrade()'s actual output -- the constant cannot drift silently when someone edits 0001. - test_legacy_backfill_skips_non_baseline_tables verifies the restricted backfill does not create a phantom table on Base.metadata, modelling a future revision that would otherwise collide on op.create_table. Doc residuals - Three-branch decision table is now consistent across bootstrap.py top docstring, engine.py comment, test module docstring, and CLAUDE.md. - Stale test anchor in blocking_io/test_persistence_engine_sqlite.py docstring now points at the real file. * fix(style): fix lint error * fix(persistence): close drift detection holes - _check_column_drift compares column type via a family equivalence allowlist ({JSON, JSONB}). Catches the wrong-type workaround `TEXT NOT NULL DEFAULT '{}'` that previously slipped through silently, while keeping Postgres JSON/JSONB dialect reflection quiet. Reflected and desired type are also echoed in every drift warning's payload for operator triage. - Extract _escape_url_for_alembic so bootstrap._alembic_safe_url and scripts/_autogen_revision share the ConfigParser % escape rule instead of duplicating it. - backend/README.md: add `make migrate-rev MSG=...` to Commands and a Schema Migrations section per the repo's README/CLAUDE.md sync policy. - test_base_to_dict.py: scope the test ORM class to an isolated MetaData so the create_all-vs-alembic parity test (added in the previous commit) is not polluted by the phantom table on the full pytest session.
292 lines
16 KiB
Python
292 lines
16 KiB
Python
"""baseline -- chain root for DeerFlow application schema.
|
|
|
|
Revision ID: 0001_baseline
|
|
Revises:
|
|
Create Date: 2026-06-22
|
|
|
|
Role of this revision
|
|
=====================
|
|
|
|
This revision encodes the schema that ``Base.metadata.create_all`` produces for
|
|
every DeerFlow-owned table at the point alembic was wired in. Under the hybrid
|
|
bootstrap strategy (``deerflow.persistence.bootstrap.bootstrap_schema``), the
|
|
``upgrade()`` here is **almost never executed**:
|
|
|
|
- Fresh DB -> ``create_all`` + ``alembic stamp head`` (no upgrade run).
|
|
- Legacy DB -> ``alembic stamp 0001_baseline`` + ``upgrade head``
|
|
(jumps directly to the next revision; baseline ``upgrade()``
|
|
is also not run, because alembic only runs revisions strictly
|
|
AFTER the stamped position).
|
|
- Versioned DB -> ``upgrade head`` (continues from whatever revision is in
|
|
``alembic_version``; baseline ``upgrade()`` only runs when
|
|
the DB happens to be at ``base``).
|
|
|
|
The baseline therefore primarily serves as a **stamp target + chain root**.
|
|
``upgrade()`` is kept faithful to ``Base.metadata`` so ``alembic upgrade base
|
|
-> head`` round-trips in test fixtures and ``downgrade()`` is provided in full
|
|
for symmetry, but production-path correctness does not depend on this
|
|
revision's DDL matching ``create_all`` byte-for-byte.
|
|
|
|
LangGraph checkpointer tables (``checkpoints``, ``checkpoint_blobs``,
|
|
``checkpoint_writes``, ``checkpoint_migrations``) are intentionally absent --
|
|
they belong to LangGraph and are excluded by ``env.py::include_object``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
|
|
import sqlalchemy as sa
|
|
from alembic import op
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = "0001_baseline"
|
|
down_revision: str | Sequence[str] | None = None
|
|
branch_labels: str | Sequence[str] | None = None
|
|
depends_on: str | Sequence[str] | None = None
|
|
|
|
|
|
def upgrade() -> None:
|
|
"""Upgrade schema."""
|
|
# ### commands auto generated by Alembic - please adjust! ###
|
|
op.create_table(
|
|
"channel_connections",
|
|
sa.Column("id", sa.String(length=64), nullable=False),
|
|
sa.Column("owner_user_id", sa.String(length=64), nullable=False),
|
|
sa.Column("provider", sa.String(length=32), nullable=False),
|
|
sa.Column("status", sa.String(length=32), nullable=False),
|
|
sa.Column("external_account_id", sa.String(length=128), nullable=False),
|
|
sa.Column("external_account_name", sa.String(length=256), nullable=True),
|
|
sa.Column("workspace_id", sa.String(length=128), nullable=False),
|
|
sa.Column("workspace_name", sa.String(length=256), nullable=True),
|
|
sa.Column("bot_user_id", sa.String(length=128), nullable=True),
|
|
sa.Column("scopes_json", sa.JSON(), nullable=False),
|
|
sa.Column("capabilities_json", sa.JSON(), nullable=False),
|
|
sa.Column("metadata_json", sa.JSON(), nullable=False),
|
|
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
|
|
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
|
|
sa.Column("last_seen_at", sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column("last_error_at", sa.DateTime(timezone=True), nullable=True),
|
|
sa.PrimaryKeyConstraint("id"),
|
|
sa.UniqueConstraint("owner_user_id", "provider", "external_account_id", "workspace_id", name="uq_channel_connection_owner_provider_identity"),
|
|
)
|
|
with op.batch_alter_table("channel_connections", schema=None) as batch_op:
|
|
batch_op.create_index("idx_channel_connections_event_lookup", ["provider", "workspace_id", "bot_user_id"], unique=False)
|
|
batch_op.create_index(batch_op.f("ix_channel_connections_owner_user_id"), ["owner_user_id"], unique=False)
|
|
batch_op.create_index(batch_op.f("ix_channel_connections_provider"), ["provider"], unique=False)
|
|
batch_op.create_index("uq_channel_connection_active_identity", ["provider", "external_account_id", "workspace_id"], unique=True, sqlite_where=sa.text("status != 'revoked'"), postgresql_where=sa.text("status != 'revoked'"))
|
|
|
|
op.create_table(
|
|
"channel_oauth_states",
|
|
sa.Column("state_hash", sa.String(length=128), nullable=False),
|
|
sa.Column("owner_user_id", sa.String(length=64), nullable=False),
|
|
sa.Column("provider", sa.String(length=32), nullable=False),
|
|
sa.Column("code_verifier_encrypted", sa.Text(), nullable=True),
|
|
sa.Column("nonce_hash", sa.String(length=128), nullable=True),
|
|
sa.Column("redirect_after", sa.Text(), nullable=True),
|
|
sa.Column("requested_scopes_json", sa.JSON(), nullable=False),
|
|
sa.Column("metadata_json", sa.JSON(), nullable=False),
|
|
sa.Column("expires_at", sa.DateTime(timezone=True), nullable=False),
|
|
sa.Column("consumed_at", sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
|
|
sa.PrimaryKeyConstraint("state_hash"),
|
|
)
|
|
with op.batch_alter_table("channel_oauth_states", schema=None) as batch_op:
|
|
batch_op.create_index(batch_op.f("ix_channel_oauth_states_owner_user_id"), ["owner_user_id"], unique=False)
|
|
batch_op.create_index(batch_op.f("ix_channel_oauth_states_provider"), ["provider"], unique=False)
|
|
|
|
op.create_table(
|
|
"feedback",
|
|
sa.Column("feedback_id", sa.String(length=64), nullable=False),
|
|
sa.Column("run_id", sa.String(length=64), nullable=False),
|
|
sa.Column("thread_id", sa.String(length=64), nullable=False),
|
|
sa.Column("user_id", sa.String(length=64), nullable=True),
|
|
sa.Column("message_id", sa.String(length=64), nullable=True),
|
|
sa.Column("rating", sa.Integer(), nullable=False),
|
|
sa.Column("comment", sa.Text(), nullable=True),
|
|
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
|
|
sa.PrimaryKeyConstraint("feedback_id"),
|
|
sa.UniqueConstraint("thread_id", "run_id", "user_id", name="uq_feedback_thread_run_user"),
|
|
)
|
|
with op.batch_alter_table("feedback", schema=None) as batch_op:
|
|
batch_op.create_index(batch_op.f("ix_feedback_run_id"), ["run_id"], unique=False)
|
|
batch_op.create_index(batch_op.f("ix_feedback_thread_id"), ["thread_id"], unique=False)
|
|
batch_op.create_index(batch_op.f("ix_feedback_user_id"), ["user_id"], unique=False)
|
|
|
|
op.create_table(
|
|
"run_events",
|
|
sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
|
|
sa.Column("thread_id", sa.String(length=64), nullable=False),
|
|
sa.Column("run_id", sa.String(length=64), nullable=False),
|
|
sa.Column("user_id", sa.String(length=64), nullable=True),
|
|
sa.Column("event_type", sa.String(length=32), nullable=False),
|
|
sa.Column("category", sa.String(length=16), nullable=False),
|
|
sa.Column("content", sa.Text(), nullable=False),
|
|
sa.Column("event_metadata", sa.JSON(), nullable=False),
|
|
sa.Column("seq", sa.Integer(), nullable=False),
|
|
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
|
|
sa.PrimaryKeyConstraint("id"),
|
|
sa.UniqueConstraint("thread_id", "seq", name="uq_events_thread_seq"),
|
|
)
|
|
with op.batch_alter_table("run_events", schema=None) as batch_op:
|
|
batch_op.create_index("ix_events_run", ["thread_id", "run_id", "seq"], unique=False)
|
|
batch_op.create_index("ix_events_thread_cat_seq", ["thread_id", "category", "seq"], unique=False)
|
|
batch_op.create_index(batch_op.f("ix_run_events_user_id"), ["user_id"], unique=False)
|
|
|
|
op.create_table(
|
|
"runs",
|
|
sa.Column("run_id", sa.String(length=64), nullable=False),
|
|
sa.Column("thread_id", sa.String(length=64), nullable=False),
|
|
sa.Column("assistant_id", sa.String(length=128), nullable=True),
|
|
sa.Column("user_id", sa.String(length=64), nullable=True),
|
|
sa.Column("status", sa.String(length=20), nullable=False),
|
|
sa.Column("model_name", sa.String(length=128), nullable=True),
|
|
sa.Column("multitask_strategy", sa.String(length=20), nullable=False),
|
|
sa.Column("metadata_json", sa.JSON(), nullable=False),
|
|
sa.Column("kwargs_json", sa.JSON(), nullable=False),
|
|
sa.Column("error", sa.Text(), nullable=True),
|
|
sa.Column("message_count", sa.Integer(), nullable=False),
|
|
sa.Column("first_human_message", sa.Text(), nullable=True),
|
|
sa.Column("last_ai_message", sa.Text(), nullable=True),
|
|
sa.Column("total_input_tokens", sa.Integer(), nullable=False),
|
|
sa.Column("total_output_tokens", sa.Integer(), nullable=False),
|
|
sa.Column("total_tokens", sa.Integer(), nullable=False),
|
|
sa.Column("llm_call_count", sa.Integer(), nullable=False),
|
|
sa.Column("lead_agent_tokens", sa.Integer(), nullable=False),
|
|
sa.Column("subagent_tokens", sa.Integer(), nullable=False),
|
|
sa.Column("middleware_tokens", sa.Integer(), nullable=False),
|
|
sa.Column("token_usage_by_model", sa.JSON(), nullable=False, server_default=sa.text("'{}'")),
|
|
sa.Column("follow_up_to_run_id", sa.String(length=64), nullable=True),
|
|
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
|
|
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
|
|
sa.PrimaryKeyConstraint("run_id"),
|
|
)
|
|
with op.batch_alter_table("runs", schema=None) as batch_op:
|
|
batch_op.create_index(batch_op.f("ix_runs_thread_id"), ["thread_id"], unique=False)
|
|
batch_op.create_index("ix_runs_thread_status", ["thread_id", "status"], unique=False)
|
|
batch_op.create_index(batch_op.f("ix_runs_user_id"), ["user_id"], unique=False)
|
|
|
|
op.create_table(
|
|
"threads_meta",
|
|
sa.Column("thread_id", sa.String(length=64), nullable=False),
|
|
sa.Column("assistant_id", sa.String(length=128), nullable=True),
|
|
sa.Column("user_id", sa.String(length=64), nullable=True),
|
|
sa.Column("display_name", sa.String(length=256), nullable=True),
|
|
sa.Column("status", sa.String(length=20), nullable=False),
|
|
sa.Column("metadata_json", sa.JSON(), nullable=False),
|
|
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
|
|
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
|
|
sa.PrimaryKeyConstraint("thread_id"),
|
|
)
|
|
with op.batch_alter_table("threads_meta", schema=None) as batch_op:
|
|
batch_op.create_index(batch_op.f("ix_threads_meta_assistant_id"), ["assistant_id"], unique=False)
|
|
batch_op.create_index(batch_op.f("ix_threads_meta_user_id"), ["user_id"], unique=False)
|
|
|
|
op.create_table(
|
|
"users",
|
|
sa.Column("id", sa.String(length=36), nullable=False),
|
|
sa.Column("email", sa.String(length=320), nullable=False),
|
|
sa.Column("password_hash", sa.String(length=128), nullable=True),
|
|
sa.Column("system_role", sa.String(length=16), nullable=False),
|
|
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
|
|
sa.Column("oauth_provider", sa.String(length=32), nullable=True),
|
|
sa.Column("oauth_id", sa.String(length=128), nullable=True),
|
|
sa.Column("needs_setup", sa.Boolean(), nullable=False),
|
|
sa.Column("token_version", sa.Integer(), nullable=False),
|
|
sa.PrimaryKeyConstraint("id"),
|
|
)
|
|
with op.batch_alter_table("users", schema=None) as batch_op:
|
|
batch_op.create_index("idx_users_oauth_identity", ["oauth_provider", "oauth_id"], unique=True, sqlite_where=sa.text("oauth_provider IS NOT NULL AND oauth_id IS NOT NULL"))
|
|
batch_op.create_index(batch_op.f("ix_users_email"), ["email"], unique=True)
|
|
|
|
op.create_table(
|
|
"channel_conversations",
|
|
sa.Column("id", sa.String(length=64), nullable=False),
|
|
sa.Column("connection_id", sa.String(length=64), nullable=False),
|
|
sa.Column("owner_user_id", sa.String(length=64), nullable=False),
|
|
sa.Column("provider", sa.String(length=32), nullable=False),
|
|
sa.Column("external_conversation_id", sa.String(length=128), nullable=False),
|
|
sa.Column("external_topic_id", sa.String(length=128), nullable=False),
|
|
sa.Column("thread_id", sa.String(length=64), nullable=False),
|
|
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
|
|
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
|
|
sa.ForeignKeyConstraint(["connection_id"], ["channel_connections.id"], ondelete="CASCADE"),
|
|
sa.PrimaryKeyConstraint("id"),
|
|
sa.UniqueConstraint("connection_id", "external_conversation_id", "external_topic_id", name="uq_channel_conversation_connection_external"),
|
|
)
|
|
with op.batch_alter_table("channel_conversations", schema=None) as batch_op:
|
|
batch_op.create_index(batch_op.f("ix_channel_conversations_connection_id"), ["connection_id"], unique=False)
|
|
batch_op.create_index(batch_op.f("ix_channel_conversations_owner_user_id"), ["owner_user_id"], unique=False)
|
|
batch_op.create_index(batch_op.f("ix_channel_conversations_provider"), ["provider"], unique=False)
|
|
batch_op.create_index(batch_op.f("ix_channel_conversations_thread_id"), ["thread_id"], unique=False)
|
|
|
|
op.create_table(
|
|
"channel_credentials",
|
|
sa.Column("connection_id", sa.String(length=64), nullable=False),
|
|
sa.Column("encrypted_access_token", sa.Text(), nullable=True),
|
|
sa.Column("encrypted_refresh_token", sa.Text(), nullable=True),
|
|
sa.Column("token_type", sa.String(length=32), nullable=True),
|
|
sa.Column("expires_at", sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column("refresh_expires_at", sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column("encrypted_extra_json", sa.Text(), nullable=True),
|
|
sa.Column("version", sa.Integer(), nullable=False),
|
|
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
|
|
sa.ForeignKeyConstraint(["connection_id"], ["channel_connections.id"], ondelete="CASCADE"),
|
|
sa.PrimaryKeyConstraint("connection_id"),
|
|
)
|
|
# ### end Alembic commands ###
|
|
|
|
|
|
def downgrade() -> None:
|
|
"""Downgrade schema."""
|
|
# ### commands auto generated by Alembic - please adjust! ###
|
|
op.drop_table("channel_credentials")
|
|
with op.batch_alter_table("channel_conversations", schema=None) as batch_op:
|
|
batch_op.drop_index(batch_op.f("ix_channel_conversations_thread_id"))
|
|
batch_op.drop_index(batch_op.f("ix_channel_conversations_provider"))
|
|
batch_op.drop_index(batch_op.f("ix_channel_conversations_owner_user_id"))
|
|
batch_op.drop_index(batch_op.f("ix_channel_conversations_connection_id"))
|
|
|
|
op.drop_table("channel_conversations")
|
|
with op.batch_alter_table("users", schema=None) as batch_op:
|
|
batch_op.drop_index(batch_op.f("ix_users_email"))
|
|
batch_op.drop_index("idx_users_oauth_identity", sqlite_where=sa.text("oauth_provider IS NOT NULL AND oauth_id IS NOT NULL"))
|
|
|
|
op.drop_table("users")
|
|
with op.batch_alter_table("threads_meta", schema=None) as batch_op:
|
|
batch_op.drop_index(batch_op.f("ix_threads_meta_user_id"))
|
|
batch_op.drop_index(batch_op.f("ix_threads_meta_assistant_id"))
|
|
|
|
op.drop_table("threads_meta")
|
|
with op.batch_alter_table("runs", schema=None) as batch_op:
|
|
batch_op.drop_index(batch_op.f("ix_runs_user_id"))
|
|
batch_op.drop_index("ix_runs_thread_status")
|
|
batch_op.drop_index(batch_op.f("ix_runs_thread_id"))
|
|
|
|
op.drop_table("runs")
|
|
with op.batch_alter_table("run_events", schema=None) as batch_op:
|
|
batch_op.drop_index(batch_op.f("ix_run_events_user_id"))
|
|
batch_op.drop_index("ix_events_thread_cat_seq")
|
|
batch_op.drop_index("ix_events_run")
|
|
|
|
op.drop_table("run_events")
|
|
with op.batch_alter_table("feedback", schema=None) as batch_op:
|
|
batch_op.drop_index(batch_op.f("ix_feedback_user_id"))
|
|
batch_op.drop_index(batch_op.f("ix_feedback_thread_id"))
|
|
batch_op.drop_index(batch_op.f("ix_feedback_run_id"))
|
|
|
|
op.drop_table("feedback")
|
|
with op.batch_alter_table("channel_oauth_states", schema=None) as batch_op:
|
|
batch_op.drop_index(batch_op.f("ix_channel_oauth_states_provider"))
|
|
batch_op.drop_index(batch_op.f("ix_channel_oauth_states_owner_user_id"))
|
|
|
|
op.drop_table("channel_oauth_states")
|
|
with op.batch_alter_table("channel_connections", schema=None) as batch_op:
|
|
batch_op.drop_index("uq_channel_connection_active_identity", sqlite_where=sa.text("status != 'revoked'"), postgresql_where=sa.text("status != 'revoked'"))
|
|
batch_op.drop_index(batch_op.f("ix_channel_connections_provider"))
|
|
batch_op.drop_index(batch_op.f("ix_channel_connections_owner_user_id"))
|
|
batch_op.drop_index("idx_channel_connections_event_lookup")
|
|
|
|
op.drop_table("channel_connections")
|
|
# ### end Alembic commands ###
|