deer-flow/backend/tests/test_migration_0029_scheduler_agent_tasks.py

"""Additive scheduler upgrade and meaningful unmet downgrade compatibility."""

from __future__ import annotations

import asyncio
import json
import os
import uuid
from datetime import UTC, datetime
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit

import pytest
import sqlalchemy as sa
from alembic import command
from alembic.autogenerate import compare_metadata
from alembic.migration import MigrationContext
from alembic.script import ScriptDirectory
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

from deerflow.persistence import bootstrap
from deerflow.persistence.base import Base
from deerflow.persistence.postgres_schema import build_asyncpg_connect_args
from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository

REVISION = "0029_scheduler_agent_tasks"
NEXT = "0030_notification_claim_tokens"
PREVIOUS = "0028_parked_attempts"
TASK_FIELDS = {"origin_thread_id", "goal_objective", "max_runs", "end_at", "standing_notes"}
OCCURRENCE_FIELDS = {"goal_objective", "goal_verdict", "stop_requested_run_id"}
pytestmark = pytest.mark.asyncio


async def test_0029_remains_in_the_single_migration_chain():
    script = ScriptDirectory(str(bootstrap._MIGRATIONS_DIR))
    assert len(script.get_heads()) == 1
    assert REVISION in {revision.revision for revision in script.walk_revisions()}
    assert script.get_revision(REVISION).down_revision == PREVIOUS
    assert script.get_revision(NEXT).down_revision == REVISION
    assert len(REVISION) <= 32


@pytest.mark.parametrize("migration_backend", ["sqlite", "postgres"])
async def test_0029_preserves_legacy_fields_and_downgrades_unmet(tmp_path, migration_backend):
    schema = None
    if migration_backend == "postgres":
        uri = os.environ.get("TEST_POSTGRES_URI")
        if not uri:
            pytest.skip("requires TEST_POSTGRES_URI (real Postgres scheduler migration)")
        parts = urlsplit(uri)
        scheme = "postgresql+asyncpg" if parts.scheme in {"postgres", "postgresql"} else parts.scheme
        query = urlencode([(key, value) for key, value in parse_qsl(parts.query, keep_blank_values=True) if key not in {"sslmode", "channel_binding"}])
        uri = urlunsplit(parts._replace(scheme=scheme, query=query))
        schema = f"scheduler_migration_{uuid.uuid4().hex}"
        engine = create_async_engine(uri, connect_args=build_asyncpg_connect_args(schema))
    else:
        engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'migration.db'}")
    cfg = bootstrap._get_alembic_config(engine, postgres_schema=schema or "")
    try:
        if schema:
            async with engine.begin() as conn:
                await conn.execute(sa.text(f'CREATE SCHEMA "{schema}"'))
        await asyncio.to_thread(command.upgrade, cfg, PREVIOUS)
        async with engine.begin() as conn:
            await conn.execute(
                sa.text(
                    "INSERT INTO scheduled_tasks (id,user_id,context_mode,title,prompt,schedule_type,schedule_spec,timezone,status,overlap_policy,run_count,created_at,updated_at) "
                    "VALUES ('legacy','owner','fresh_thread_per_run','Legacy','Prompt','cron','{}','UTC','enabled','enqueue',7,CURRENT_TIMESTAMP,CURRENT_TIMESTAMP)"
                )
            )
            await conn.execute(sa.text("INSERT INTO scheduled_task_runs (id,task_id,thread_id,scheduled_for,trigger,status,created_at) VALUES ('old-occurrence','legacy','thread',CURRENT_TIMESTAMP,'scheduled','success',CURRENT_TIMESTAMP)"))
        await asyncio.to_thread(command.upgrade, cfg, REVISION)
        async with engine.connect() as conn:
            columns = await conn.run_sync(lambda sync: {table: {col["name"]: col for col in sa.inspect(sync).get_columns(table)} for table in ("runs", "scheduled_tasks", "scheduled_task_runs")})
        for table, fields in [("runs", {"goal_verdict"}), ("scheduled_tasks", TASK_FIELDS), ("scheduled_task_runs", OCCURRENCE_FIELDS)]:
            assert all(columns[table][field]["nullable"] and columns[table][field]["default"] is None for field in fields)
        # Compare with the ORM at the current head: later revisions add
        # scheduled_tasks columns (0031) that the model already declares.
        await asyncio.to_thread(command.upgrade, cfg, "head")
        async with engine.connect() as conn:

            def scheduler_schema_diff(sync):
                targets = {"runs", "scheduled_tasks", "scheduled_task_runs"}

                def include_object(obj, name, type_, reflected, compare_to):
                    return name in targets if type_ == "table" else obj.table.name in targets

                def compare_default(context, inspected_column, metadata_column, inspected_default, metadata_default, rendered_default):
                    if isinstance(metadata_column.type, sa.JSON):
                        # PostgreSQL JSON has no equality operator. Compare
                        # decoded defaults without executing JSON '=' in SQL.
                        def decoded(value):
                            return None if value is None else json.loads(value.rsplit("::", 1)[0].strip().strip("'"))

                        return decoded(inspected_default) != decoded(rendered_default)
                    return None

                context = MigrationContext.configure(sync, opts={"include_object": include_object, "compare_type": True, "compare_server_default": compare_default})
                return compare_metadata(context, Base.metadata)

            assert await conn.run_sync(scheduler_schema_diff) == []
        repo = ScheduledTaskRepository(async_sessionmaker(engine, expire_on_commit=False))
        legacy = await repo.get("legacy", user_id="owner")
        assert legacy["run_count"] == 7
        assert all(legacy[field] is None for field in TASK_FIELDS)
        await repo.create(
            task_id="new",
            user_id="owner",
            thread_id=None,
            context_mode="fresh_thread_per_run",
            assistant_id=None,
            title="new",
            prompt="test",
            schedule_type="interval",
            schedule_spec={"every_seconds": 3600},
            timezone="UTC",
            next_run_at=datetime.now(UTC),
            origin_thread_id="origin",
            goal_objective="deliver",
            max_runs=2,
            standing_notes=["develop branch"],
        )
        async with engine.begin() as conn:
            await conn.execute(sa.text("UPDATE scheduled_task_runs SET status = 'unmet', error = 'no_verdict' WHERE id = 'old-occurrence'"))
        await asyncio.to_thread(command.downgrade, cfg, PREVIOUS)
        async with engine.connect() as conn:
            row = (await conn.execute(sa.text("SELECT status,error FROM scheduled_task_runs WHERE id='old-occurrence'"))).one()
            assert row == ("failed", "no_verdict")
            remaining = await conn.run_sync(lambda sync: {col["name"] for col in sa.inspect(sync).get_columns("scheduled_tasks")})
            assert TASK_FIELDS.isdisjoint(remaining)
            assert (await conn.execute(sa.text("SELECT run_count FROM scheduled_tasks WHERE id='legacy'"))).scalar_one() == 7
        # The ORM model carries later scheduled_tasks columns (0031), so read
        # it back at the current head.
        await asyncio.to_thread(command.upgrade, cfg, "head")
        assert (await repo.get("new", user_id="owner"))["origin_thread_id"] is None
    finally:
        if schema:
            async with engine.begin() as conn:
                await conn.execute(sa.text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
        await engine.dispose()