Skip to content

fix(DB): check scheduled_task_runs independently in 0003 migration guard #4198

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
DaoyuanLi2816 wants to merge 1 commit into
base: main
Choose a base branch
from
+281 −55
Open

fix(DB): check scheduled_task_runs independently in 0003 migration guard #4198

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
120 changes: 65 additions & 55 deletions backend/packages/harness/deerflow/persistence/migrations/versions/0003_scheduled_tasks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,62 +21,72 @@
def upgrade() -> None:
bind = op.get_bind()
inspector = sa.inspect(bind)
if inspector.has_table("scheduled_tasks"):
# Idempotent: a DB whose full-metadata create_all already provisioned
# both scheduled-task tables (e.g. legacy test seeds) must not have them
# re-created here.
return
op.create_table(
"scheduled_tasks",
sa.Column("id", sa.String(length=64), nullable=False),
sa.Column("user_id", sa.String(length=64), nullable=False),
sa.Column("thread_id", sa.String(length=64), nullable=True),
sa.Column("context_mode", sa.String(length=32), nullable=False),
sa.Column("assistant_id", sa.String(length=128), nullable=True),
sa.Column("title", sa.String(length=255), nullable=False),
sa.Column("prompt", sa.Text(), nullable=False),
sa.Column("schedule_type", sa.String(length=16), nullable=False),
sa.Column("schedule_spec", sa.JSON(), nullable=False),
sa.Column("timezone", sa.String(length=64), nullable=False),
sa.Column("status", sa.String(length=16), nullable=False),
sa.Column("overlap_policy", sa.String(length=16), nullable=False),
sa.Column("next_run_at", sa.DateTime(timezone=True), nullable=True),
sa.Column("last_run_at", sa.DateTime(timezone=True), nullable=True),
sa.Column("last_run_id", sa.String(length=64), nullable=True),
sa.Column("last_thread_id", sa.String(length=64), nullable=True),
sa.Column("last_error", sa.Text(), nullable=True),
sa.Column("lease_owner", sa.String(length=128), nullable=True),
sa.Column("lease_expires_at", sa.DateTime(timezone=True), nullable=True),
sa.Column("run_count", sa.Integer(), nullable=False),
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
sa.PrimaryKeyConstraint("id"),
)
with op.batch_alter_table("scheduled_tasks", schema=None) as batch_op:
batch_op.create_index("ix_scheduled_tasks_user_id", ["user_id"], unique=False)
batch_op.create_index("ix_scheduled_tasks_thread_id", ["thread_id"], unique=False)
batch_op.create_index("ix_scheduled_tasks_status", ["status"], unique=False)
batch_op.create_index("ix_scheduled_tasks_next_run_at", ["next_run_at"], unique=False)

op.create_table(
"scheduled_task_runs",
sa.Column("id", sa.String(length=64), nullable=False),
sa.Column("task_id", sa.String(length=64), nullable=False),
sa.Column("thread_id", sa.String(length=64), nullable=False),
sa.Column("run_id", sa.String(length=64), nullable=True),
sa.Column("scheduled_for", sa.DateTime(timezone=True), nullable=False),
sa.Column("trigger", sa.String(length=16), nullable=False),
sa.Column("status", sa.String(length=16), nullable=False),
sa.Column("error", sa.Text(), nullable=True),
sa.Column("started_at", sa.DateTime(timezone=True), nullable=True),
sa.Column("finished_at", sa.DateTime(timezone=True), nullable=True),
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
sa.PrimaryKeyConstraint("id"),
)
with op.batch_alter_table("scheduled_task_runs", schema=None) as batch_op:
batch_op.create_index("ix_scheduled_task_runs_task_id", ["task_id"], unique=False)
batch_op.create_index("ix_scheduled_task_runs_thread_id", ["thread_id"], unique=False)
batch_op.create_index("ix_scheduled_task_runs_status", ["status"], unique=False)
# Idempotent per-table checks (mirrors the per-object pattern in
# 0004_run_ownership.py's index creation): a DB whose full-metadata
# create_all already provisioned one or both scheduled-task tables (e.g.
# legacy test seeds, a restore, or manual DB surgery) must not have an
# already-present table re-created here. The two tables are checked
# independently -- rather than gating the whole revision on just
# ``scheduled_tasks`` -- so a DB that atypically has one table but not
# the other still gets the missing one created instead of the revision
# silently skipping it.
has_scheduled_tasks = inspector.has_table("scheduled_tasks")
has_scheduled_task_runs = inspector.has_table("scheduled_task_runs")

if not has_scheduled_tasks:
op.create_table(
"scheduled_tasks",
sa.Column("id", sa.String(length=64), nullable=False),
sa.Column("user_id", sa.String(length=64), nullable=False),
sa.Column("thread_id", sa.String(length=64), nullable=True),
sa.Column("context_mode", sa.String(length=32), nullable=False),
sa.Column("assistant_id", sa.String(length=128), nullable=True),
sa.Column("title", sa.String(length=255), nullable=False),
sa.Column("prompt", sa.Text(), nullable=False),
sa.Column("schedule_type", sa.String(length=16), nullable=False),
sa.Column("schedule_spec", sa.JSON(), nullable=False),
sa.Column("timezone", sa.String(length=64), nullable=False),
sa.Column("status", sa.String(length=16), nullable=False),
sa.Column("overlap_policy", sa.String(length=16), nullable=False),
sa.Column("next_run_at", sa.DateTime(timezone=True), nullable=True),
sa.Column("last_run_at", sa.DateTime(timezone=True), nullable=True),
sa.Column("last_run_id", sa.String(length=64), nullable=True),
sa.Column("last_thread_id", sa.String(length=64), nullable=True),
sa.Column("last_error", sa.Text(), nullable=True),
sa.Column("lease_owner", sa.String(length=128), nullable=True),
sa.Column("lease_expires_at", sa.DateTime(timezone=True), nullable=True),
sa.Column("run_count", sa.Integer(), nullable=False),
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
sa.PrimaryKeyConstraint("id"),
)
with op.batch_alter_table("scheduled_tasks", schema=None) as batch_op:
batch_op.create_index("ix_scheduled_tasks_user_id", ["user_id"], unique=False)
batch_op.create_index("ix_scheduled_tasks_thread_id", ["thread_id"], unique=False)
batch_op.create_index("ix_scheduled_tasks_status", ["status"], unique=False)
batch_op.create_index("ix_scheduled_tasks_next_run_at", ["next_run_at"], unique=False)

if not has_scheduled_task_runs:
op.create_table(
"scheduled_task_runs",
sa.Column("id", sa.String(length=64), nullable=False),
sa.Column("task_id", sa.String(length=64), nullable=False),
sa.Column("thread_id", sa.String(length=64), nullable=False),
sa.Column("run_id", sa.String(length=64), nullable=True),
sa.Column("scheduled_for", sa.DateTime(timezone=True), nullable=False),
sa.Column("trigger", sa.String(length=16), nullable=False),
sa.Column("status", sa.String(length=16), nullable=False),
sa.Column("error", sa.Text(), nullable=True),
sa.Column("started_at", sa.DateTime(timezone=True), nullable=True),
sa.Column("finished_at", sa.DateTime(timezone=True), nullable=True),
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
sa.PrimaryKeyConstraint("id"),
)
with op.batch_alter_table("scheduled_task_runs", schema=None) as batch_op:
batch_op.create_index("ix_scheduled_task_runs_task_id", ["task_id"], unique=False)
batch_op.create_index("ix_scheduled_task_runs_thread_id", ["thread_id"], unique=False)
batch_op.create_index("ix_scheduled_task_runs_status", ["status"], unique=False)


def downgrade() -> None:
Expand Down
216 changes: 216 additions & 0 deletions backend/tests/test_migration_0003_scheduled_task_runs_guard.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,216 @@
"""Regression test for migration ``0003_scheduled_tasks``'s idempotency guard.

Revision ``0003_scheduled_tasks`` creates BOTH ``scheduled_tasks`` and
``scheduled_task_runs`` in one ``upgrade()``, but its re-entry guard used to
check only the first table:

if inspector.has_table("scheduled_tasks"):
return

On a DB that atypically already has ``scheduled_tasks`` but not its sibling
``scheduled_task_runs`` -- a partial restore, manual DB surgery, or legacy
provisioning that only ever created the first table -- the guard returns
early and silently skips creating ``scheduled_task_runs``. ``alembic_version``
still advances past ``0003`` (and on to head), so there is no error and no
retry path: the gap is permanent once triggered, not transient. The next code
path that touches ``scheduled_task_runs`` (e.g.
``ScheduledTaskRunRepository.count_active_runs()``) then 500s with
``OperationalError: no such table: scheduled_task_runs``.

Reaching this requires the atypical prior state described above. The ordinary
empty-DB and legacy-DB bootstrap branches (``persistence/bootstrap.py``) both
provision ``scheduled_tasks`` and ``scheduled_task_runs`` atomically via
``Base.metadata.create_all``, so a normal fresh install or upgrade never hits
the asymmetric state this test seeds.

End-to-end shape (mirrors ``test_migration_0004_run_ownership_dedupe.py``):

1. Hand-build a SQLite DB at revision ``0002_runs_token_usage`` where
``scheduled_tasks`` already exists but ``scheduled_task_runs`` does not.
2. Run ``init_engine`` (the FastAPI lifespan entry point), which routes
through ``bootstrap_schema`` -> the versioned branch -> ``alembic upgrade
head`` -> ``0003.upgrade()`` (the revision under test) -> ``0004.upgrade()``.
3. Verify ``scheduled_task_runs`` now exists, ``alembic_version`` reached
head, and the real repository method succeeds instead of raising.

Two companion tests pin the cases the fix must not regress: the ordinary path
(neither table exists) still creates both, and the fully-idempotent path
(both already exist, e.g. via the hybrid bootstrap's atomic ``create_all``)
still no-ops cleanly with no ``table ... already exists`` error.
"""

from __future__ import annotations

import sqlite3
from pathlib import Path

import pytest
import sqlalchemy as sa

import deerflow.persistence.models # noqa: F401 -- registers ORM models
from deerflow.persistence.base import Base
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
from deerflow.persistence.scheduled_task_runs.sql import ScheduledTaskRunRepository

pytestmark = pytest.mark.asyncio

HEAD = "0004_run_ownership"
PRE_0003_REVISION = "0002_runs_token_usage"


def _stamp_alembic_version(conn: sa.Connection, revision: str) -> None:
conn.execute(sa.text("CREATE TABLE IF NOT EXISTS alembic_version (version_num VARCHAR(32) NOT NULL)"))
conn.execute(sa.text("DELETE FROM alembic_version"))
conn.execute(sa.text("INSERT INTO alembic_version (version_num) VALUES (:v)"), {"v": revision})


def _table_names(db_path: Path) -> set[str]:
with sqlite3.connect(db_path) as raw:
return {row[0] for row in raw.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall()}


def _alembic_version(db_path: Path) -> str | None:
with sqlite3.connect(db_path) as raw:
row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
return row[0] if row else None


def _seed_pre_0003_with_only_scheduled_tasks(db_path: Path) -> None:
"""Build a DB at revision 0002 where ``scheduled_tasks`` exists but
``scheduled_task_runs`` does not -- the atypical prior state that trips
the 0003 guard gap.

``Base.metadata.create_all`` renders the full *current* schema (every
table across every revision, since the ORM models are the living source
of truth), so we build that and then drop just ``scheduled_task_runs``,
landing on exactly the asymmetric state the bug targets. Revision 0004's
own idempotent checks (``safe_add_column`` / the ``ix_runs_lease`` guard)
no-op harmlessly when later replayed against a create_all-rendered
``runs`` table, so leaving ``runs`` at its full current shape does not
interfere with reproducing the 0003 bug.
"""
db_path.parent.mkdir(parents=True, exist_ok=True)
sync_engine = sa.create_engine(f"sqlite:///{db_path.as_posix()}")
try:
Base.metadata.create_all(sync_engine)
with sync_engine.begin() as conn:
conn.execute(sa.text("DROP TABLE scheduled_task_runs"))
# Stamp at 0002 so bootstrap takes the versioned branch and
# replays 0003 (the revision under test) then 0004 during
# ``alembic upgrade head``.
_stamp_alembic_version(conn, PRE_0003_REVISION)
finally:
sync_engine.dispose()


def _seed_pre_0003_with_neither_scheduled_table(db_path: Path) -> None:
"""Build an ordinary DB at revision 0002: neither scheduled-task table
exists yet, matching a real deployment that has never run 0003."""
db_path.parent.mkdir(parents=True, exist_ok=True)
sync_engine = sa.create_engine(f"sqlite:///{db_path.as_posix()}")
try:
Base.metadata.create_all(sync_engine)
with sync_engine.begin() as conn:
conn.execute(sa.text("DROP TABLE scheduled_task_runs"))
conn.execute(sa.text("DROP TABLE scheduled_tasks"))
_stamp_alembic_version(conn, PRE_0003_REVISION)
finally:
sync_engine.dispose()


def _seed_pre_0003_with_both_scheduled_tables(db_path: Path) -> None:
"""Build a DB at revision 0002 where BOTH scheduled-task tables already
exist -- the fully-idempotent case (e.g. the hybrid bootstrap's legacy
branch already provisioned them atomically via ``create_all``)."""
db_path.parent.mkdir(parents=True, exist_ok=True)
sync_engine = sa.create_engine(f"sqlite:///{db_path.as_posix()}")
try:
Base.metadata.create_all(sync_engine)
with sync_engine.begin() as conn:
_stamp_alembic_version(conn, PRE_0003_REVISION)
finally:
sync_engine.dispose()


async def test_scheduled_task_runs_created_when_only_scheduled_tasks_preexists(tmp_path: Path) -> None:
db_path = tmp_path / "partial.db"
_seed_pre_0003_with_only_scheduled_tasks(db_path)

# Sanity: confirm we landed in the atypical pre-fix shape before
# init_engine touches the file.
tables_before = _table_names(db_path)
assert "scheduled_tasks" in tables_before
assert "scheduled_task_runs" not in tables_before
assert _alembic_version(db_path) == PRE_0003_REVISION

url = f"sqlite+aiosqlite:///{db_path.as_posix()}"
await init_engine(backend="sqlite", url=url, sqlite_dir=str(tmp_path))

try:
tables_after = _table_names(db_path)

# The guard gap in 0003 checked only ``scheduled_tasks`` before
# returning early, silently skipping ``scheduled_task_runs``. Both
# tables are owned by the SAME revision, so both must exist once
# alembic_version has advanced past it.
assert "scheduled_task_runs" in tables_after, "0003_scheduled_tasks guard skipped creating scheduled_task_runs when scheduled_tasks already existed"
# alembic still advances to head regardless -- the gap is silent,
# not a raised error, which is what makes it a permanent gap rather
# than something a second "restart" upgrade-head would self-heal.
assert _alembic_version(db_path) == HEAD

# Real repository call, not just schema inspection: must not raise
# ``sqlalchemy.exc.OperationalError: no such table: scheduled_task_runs``.
sf = get_session_factory()
assert sf is not None
repo = ScheduledTaskRunRepository(sf)
count = await repo.count_active_runs()
assert count == 0
finally:
await close_engine()


async def test_neither_scheduled_table_exists_both_created_normally(tmp_path: Path) -> None:
"""Ordinary-path guard: a real DB at 0002 (neither scheduled table
present) must still get BOTH tables from a single upgrade head. The fix
must not regress the common case."""
db_path = tmp_path / "neither.db"
_seed_pre_0003_with_neither_scheduled_table(db_path)

tables_before = _table_names(db_path)
assert "scheduled_tasks" not in tables_before
assert "scheduled_task_runs" not in tables_before

url = f"sqlite+aiosqlite:///{db_path.as_posix()}"
await init_engine(backend="sqlite", url=url, sqlite_dir=str(tmp_path))
try:
tables_after = _table_names(db_path)
assert "scheduled_tasks" in tables_after
assert "scheduled_task_runs" in tables_after
assert _alembic_version(db_path) == HEAD
finally:
await close_engine()


async def test_both_scheduled_tables_already_exist_upgrade_head_is_idempotent(tmp_path: Path) -> None:
"""Fully-idempotent case: both tables already exist (e.g. the hybrid
bootstrap's legacy branch already provisioned them atomically). Upgrade
head must no-op cleanly -- no ``OperationalError: table ... already
exists``."""
db_path = tmp_path / "both.db"
_seed_pre_0003_with_both_scheduled_tables(db_path)

tables_before = _table_names(db_path)
assert "scheduled_tasks" in tables_before
assert "scheduled_task_runs" in tables_before

url = f"sqlite+aiosqlite:///{db_path.as_posix()}"
# Must not raise (e.g. "table scheduled_tasks already exists").
await init_engine(backend="sqlite", url=url, sqlite_dir=str(tmp_path))
try:
tables_after = _table_names(db_path)
assert "scheduled_tasks" in tables_after
assert "scheduled_task_runs" in tables_after
assert _alembic_version(db_path) == HEAD
finally:
await close_engine()
Loading