deerflow-code/offline-backend-20260512/backend/scripts/migrate_maf_roundtable_phase2.py

"""Phase 2 migration for the MAF roundtable experiment.

Idempotent migration that brings a Phase-1 deployment up to Phase-2 schema:

1. Create ``maf_roundtable_seat_runs`` (MAF↔DeerFlow run mapping, survives restart).
2. Add ``source_event_id`` column to ``maf_roundtable_events`` (durable event dedup).
3. Create the ``(run_id, source_event_id)`` unique index on ``maf_roundtable_events``.
4. Create ``maf_roundtable_action_secrets`` (encrypted actual execution params, server-only).
5. Add ``turn_index`` column to ``maf_roundtable_actions`` (Phase-1 DBs may lack it).

Supports SQLite, PostgreSQL, and MySQL. Each step checks whether the object
already exists before creating it, so running the script multiple times has no
side effect. Old event rows keep ``source_event_id = NULL`` (SQL treats NULLs as
distinct in a UNIQUE constraint, so they are never deduped).

**Deployment order**: run this migration BEFORE enabling the
``maf_roundtable.enabled`` feature flag.

Usage (from ``offline-backend-20260512/backend/``)::

    # Preview (no writes)
    PYTHONPATH=. uv run python scripts/migrate_maf_roundtable_phase2.py --dry-run

    # Execute
    PYTHONPATH=. uv run python scripts/migrate_maf_roundtable_phase2.py
"""

from __future__ import annotations

import argparse
import asyncio
import logging
import sys
from typing import Any

logging.basicConfig(level=logging.INFO, format="%(message)s")
logger = logging.getLogger("migrate_maf_roundtable_phase2")


async def _migrate(dry_run: bool) -> int:
    from deerflow.config import get_app_config
    from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config

    config = get_app_config()
    await init_engine_from_config(config.database)
    sf = get_session_factory()
    if sf is None:
        logger.error("Database backend is 'memory' — nothing to migrate.")
        return 0

    engine = sf.kw["bind"]  # the async engine

    # Run introspection + DDL on the sync connection (works across all dialects).
    async with engine.begin() as conn:
        await conn.run_sync(_do_migrate_sync, dry_run=dry_run, logger=logger, changes_ref=[0])

    await close_engine()
    logger.info("Migration complete (dry_run=%s).", dry_run)
    return 0


def _do_migrate_sync(
    sync_conn: Any,
    *,
    dry_run: bool,
    logger: logging.Logger,
    changes_ref: list[int],
) -> None:
    """Run all migration steps on a sync connection (called via run_sync)."""
    from sqlalchemy import inspect, text

    insp = inspect(sync_conn)
    dialect_name = sync_conn.dialect.name

    def _exec(sql: str) -> None:
        if dry_run:
            logger.info("  [dry-run] would execute: %s", sql.strip()[:300])
            return
        sync_conn.execute(text(sql))
        changes_ref[0] += 1

    # ── 1. maf_roundtable_seat_runs ──────────────────────────────────────
    if "maf_roundtable_seat_runs" not in insp.get_table_names():
        logger.info("Creating table: maf_roundtable_seat_runs")
        for stmt in _DDL_SEAT_RUNS:
            _exec(stmt)
    else:
        logger.info("Table maf_roundtable_seat_runs already exists — skip.")

    # ── 2. source_event_id column on maf_roundtable_events ──────────────
    if "maf_roundtable_events" in insp.get_table_names():
        event_cols = {c["name"] for c in insp.get_columns("maf_roundtable_events")}
        if "source_event_id" not in event_cols:
            logger.info("Adding column source_event_id to maf_roundtable_events")
            _exec(_ddl_add_source_event_id(dialect_name))
        else:
            logger.info("Column source_event_id already exists on maf_roundtable_events — skip.")
    else:
        logger.info("Table maf_roundtable_events not present — skip source_event_id migration.")

    # ── 3. (run_id, source_event_id) unique index ───────────────────────
    index_found = False
    if "maf_roundtable_events" in insp.get_table_names():
        for idx in insp.get_indexes("maf_roundtable_events"):
            if idx.get("name") == "uq_maf_events_run_source":
                index_found = True
                break
    if not index_found and "maf_roundtable_events" in insp.get_table_names():
        logger.info("Creating unique index uq_maf_events_run_source on maf_roundtable_events")
        _exec(_DDL_EVENTS_SOURCE_INDEX)
    else:
        logger.info("Index uq_maf_events_run_source already exists or events table absent — skip.")

    # ── 4. maf_roundtable_action_secrets ─────────────────────────────────
    if "maf_roundtable_action_secrets" not in insp.get_table_names():
        logger.info("Creating table: maf_roundtable_action_secrets")
        for stmt in _DDL_ACTION_SECRETS:
            _exec(stmt)
    else:
        logger.info("Table maf_roundtable_action_secrets already exists — skip.")

    # ── 5. turn_index on maf_roundtable_actions (Phase-1 deployments) ───
    action_cols = (
        {c["name"] for c in insp.get_columns("maf_roundtable_actions")}
        if "maf_roundtable_actions" in insp.get_table_names()
        else set()
    )
    if "maf_roundtable_actions" in insp.get_table_names() and "turn_index" not in action_cols:
        logger.info("Adding column turn_index to maf_roundtable_actions")
        _exec(_ddl_add_action_turn_index(dialect_name))
    elif "turn_index" in action_cols:
        logger.info("Column turn_index already exists on maf_roundtable_actions — skip.")

    logger.info("Migration steps: %d applied.", changes_ref[0])


# ── DDL templates (portable across sqlite / postgres / mysql) ──────────────
# Each DDL is a LIST of single SQL statements so ``text()`` receives one
# statement at a time — multi-statement strings fail on MySQL and several
# SQLite builds because ``text()`` does not split on ``;``.

_DDL_SEAT_RUNS: list[str] = [
    """
    CREATE TABLE maf_roundtable_seat_runs (
        id VARCHAR(64) NOT NULL PRIMARY KEY,
        maf_run_id VARCHAR(64) NOT NULL,
        agent_id VARCHAR(64) NOT NULL,
        turn_index INTEGER NOT NULL,
        deerflow_run_id VARCHAR(64),
        deerflow_thread_id VARCHAR(64),
        sandbox_id VARCHAR(64),
        workflow_checkpoint_id VARCHAR(128),
        status VARCHAR(32) DEFAULT 'pending' NOT NULL,
        started_at DATETIME,
        finished_at DATETIME,
        error TEXT,
        created_at DATETIME NOT NULL,
        updated_at DATETIME NOT NULL,
        FOREIGN KEY (maf_run_id) REFERENCES maf_roundtable_runs (id) ON DELETE CASCADE
    )
    """,
    "CREATE INDEX ix_maf_seat_runs_maf_run_id ON maf_roundtable_seat_runs (maf_run_id)",
    "CREATE INDEX ix_maf_seat_runs_agent_id ON maf_roundtable_seat_runs (agent_id)",
    "CREATE INDEX ix_maf_seat_runs_status ON maf_roundtable_seat_runs (status)",
    "CREATE UNIQUE INDEX uq_maf_seat_runs_run_agent_turn ON maf_roundtable_seat_runs (maf_run_id, agent_id, turn_index)",
]

_DDL_EVENTS_SOURCE_INDEX = "CREATE UNIQUE INDEX uq_maf_events_run_source ON maf_roundtable_events (run_id, source_event_id)"

_DDL_ACTION_SECRETS: list[str] = [
    """
    CREATE TABLE maf_roundtable_action_secrets (
        id VARCHAR(64) NOT NULL PRIMARY KEY,
        action_id VARCHAR(64) NOT NULL,
        arguments_encrypted TEXT NOT NULL,
        created_at DATETIME NOT NULL,
        FOREIGN KEY (action_id) REFERENCES maf_roundtable_actions (id) ON DELETE CASCADE
    )
    """,
    "CREATE INDEX ix_maf_action_secrets_action_id ON maf_roundtable_action_secrets (action_id)",
    "CREATE UNIQUE INDEX uq_maf_action_secrets_action ON maf_roundtable_action_secrets (action_id)",
]


def _ddl_add_source_event_id(dialect: str) -> str:
    """ALTER TABLE to add source_event_id column (dialect-specific NULL handling)."""
    col_type = "VARCHAR(191)"
    return f"ALTER TABLE maf_roundtable_events ADD COLUMN source_event_id {col_type} NULL"


def _ddl_add_action_turn_index(dialect: str) -> str:
    if dialect == "mysql":
        return "ALTER TABLE maf_roundtable_actions ADD COLUMN turn_index INT NULL"
    return "ALTER TABLE maf_roundtable_actions ADD COLUMN turn_index INTEGER NULL"


def main() -> int:
    parser = argparse.ArgumentParser(description="Phase 2 migration for MAF roundtable tables")
    parser.add_argument("--dry-run", action="store_true", help="Preview only; do not write to the database")
    args = parser.parse_args()
    return asyncio.run(_migrate(args.dry_run))


if __name__ == "__main__":
    sys.exit(main())