"""Phase 3 migration for the MAF roundtable experiment. Idempotent migration that brings a Phase-2 deployment up to Phase-3 schema: 1. Add the non-null ``execution_stage`` column to ``maf_roundtable_seat_runs`` (existing rows default to ``legacy``). 2. Replace the Phase-2 three-column unique constraint ``(maf_run_id, agent_id, turn_index)`` with the Phase-3 four-column constraint ``(maf_run_id, agent_id, turn_index, execution_stage)`` so the same business round can hold multiple distinct DeerFlow executions (planner / expert / format_repair / cross_review / finalizer / group_chat) while each stage is still unique per (run, agent, turn). The business ``turn_index`` is NEVER mutated to dodge a collision — it stays the stable business round used by the frontend, approval matching and event tracing. ``execution_stage`` is the new disambiguator. Dialect handling ---------------- - **PostgreSQL / MySQL**: ``ALTER TABLE`` to add the column, then ``DROP`` the old unique constraint/index and ``ADD`` the new one. Each step is guarded by introspection so re-runs are no-ops. - **SQLite**: a unique constraint cannot be dropped in place, so the table is rebuilt: a temp table with the new shape is created, data is copied (existing rows get ``execution_stage='legacy'``), the old table is dropped, and the temp is renamed. This is the only safe way to change a SQLite unique constraint; it is fully idempotent because the rebuild only runs when the ``execution_stage`` column is absent or the old constraint is still present. **Deployment order**: run AFTER the Phase-2 migration and BEFORE enabling the Phase-3 code path (the four-column constraint is required for the new ``execution_stage`` writes to coexist within one turn). Usage (from ``offline-backend-20260512/backend/``):: # Preview (no writes) PYTHONPATH=. uv run python scripts/migrate_maf_roundtable_phase3.py --dry-run # Execute PYTHONPATH=. uv run python scripts/migrate_maf_roundtable_phase3.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_phase3") OLD_CONSTRAINT_NAME = "uq_maf_seat_runs_run_agent_turn" NEW_CONSTRAINT_NAME = "uq_maf_seat_runs_run_agent_turn_stage" TABLE = "maf_roundtable_seat_runs" 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 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 if TABLE not in insp.get_table_names(): logger.info("Table %s does not exist — Phase-3 migration is a no-op (table will be created by the ORM with the new schema).", TABLE) return columns = {c["name"] for c in insp.get_columns(TABLE)} has_stage_col = "execution_stage" in columns has_snapshot_col = "result_snapshot" in columns # ── 1a. Add execution_stage column (additive; old rows → 'legacy') ───── if not has_stage_col: logger.info("Adding column execution_stage to %s", TABLE) _exec(_ddl_add_execution_stage(dialect_name)) else: logger.info("Column execution_stage already exists on %s — skip.", TABLE) # ── 2. Replace the unique constraint (3-tuple → 4-tuple) ─────────────── # NOTE: introspection is re-fetched after the column add so the inspector # reflects the new column within this same transaction. insp = inspect(sync_conn) old_constraint = _find_unique_constraint(insp, TABLE, expected_cols={"maf_run_id", "agent_id", "turn_index"}) new_constraint = _find_unique_constraint(insp, TABLE, expected_cols={"maf_run_id", "agent_id", "turn_index", "execution_stage"}) has_new_unique = new_constraint is not None has_stage_check = _has_check_constraint(insp, TABLE, "ck_maf_seat_runs_execution_stage") if dialect_name == "sqlite": needs_rebuild = ( not has_new_unique or not has_snapshot_col or not has_stage_check ) if needs_rebuild: logger.info("SQLite: rebuilding %s (has_new_unique=%s, has_snapshot=%s, has_check=%s).", TABLE, has_new_unique, has_snapshot_col, has_stage_check) for stmt in _sqlite_rebuild_seat_runs( source_has_stage=has_stage_col, source_has_snapshot=has_snapshot_col, ): _exec(stmt) else: logger.info("SQLite: %s already fully migrated — no-op.", TABLE) logger.info("Migration steps: %d applied.", changes_ref[0]) return # PostgreSQL / MySQL if not has_new_unique: if old_constraint is not None: old_name = old_constraint.get("name") or OLD_CONSTRAINT_NAME old_kind = old_constraint.get("kind") if old_kind == "index": logger.info("Dropping old unique INDEX %s (kind=index)", old_name) _exec(_ddl_drop_index(old_name)) else: logger.info("Dropping old unique CONSTRAINT %s (kind=constraint)", old_name) _exec(_ddl_drop_constraint(dialect_name, old_name)) else: logger.info( "Old unique object not found by introspection — migration assumes " "the four-column constraint is being created fresh." ) logger.info("Adding new four-column unique constraint %s", NEW_CONSTRAINT_NAME) _exec(_ddl_add_new_constraint(dialect_name)) if not has_snapshot_col: insp = inspect(sync_conn) cols_now = {c["name"] for c in insp.get_columns(TABLE)} if "result_snapshot" not in cols_now: logger.info("Adding column result_snapshot to %s", TABLE) _exec(_ddl_add_result_snapshot(dialect_name)) insp = inspect(sync_conn) if not _has_check_constraint(insp, TABLE, "ck_maf_seat_runs_execution_stage"): logger.info("Adding CHECK constraint ck_maf_seat_runs_execution_stage") _exec(_ddl_add_execution_stage_check(dialect_name)) logger.info("Migration steps: %d applied.", changes_ref[0]) def _find_unique_constraint(insp: Any, table: str, *, expected_cols: set[str]) -> dict[str, Any] | None: """Return the unique constraint (or unique index) matching ``expected_cols``. Looks at both table-level unique constraints AND unique indexes — Phase-2 created the constraint via ``CREATE UNIQUE INDEX`` which SQLAlchemy reports as an index, while the ORM model declares it as a UniqueConstraint (reported as a constraint). We accept either form so the migration is robust to how the table was originally created. """ for uc in insp.get_unique_constraints(table): if set(uc.get("column_names") or []) == expected_cols: return {"name": uc.get("name"), "kind": "constraint"} for idx in insp.get_indexes(table): if idx.get("unique") and set(idx.get("column_names") or []) == expected_cols: return {"name": idx.get("name"), "kind": "index"} return None def _has_check_constraint(insp: Any, table: str, name: str) -> bool: """True if a CHECK constraint with the given name exists on ``table``. SQLAlchemy's inspector exposes CHECK constraints via ``get_check_constraints`` on PostgreSQL / MySQL / modern SQLite. We match by name; if the dialect cannot list checks we conservatively return False so the (idempotent) ADD runs once — a duplicate-name error there is harmless and caught by the ``IF NOT EXISTS``-style guard in the caller's re-introspection. """ getter = getattr(insp, "get_check_constraints", None) if getter is None: return False try: for ck in getter(table): if ck.get("name") == name: return True except Exception: # noqa: BLE001 return False return False # ── DDL templates (portable) ──────────────────────────────────────────────── def _ddl_add_execution_stage(dialect: str) -> str: """Add the execution_stage column, back-filling existing rows to 'legacy'. Adding NOT NULL with a server default is safe on all three dialects: the default populates existing rows during the ALTER, then stays for new ones. """ if dialect == "mysql": return ( "ALTER TABLE maf_roundtable_seat_runs " "ADD COLUMN execution_stage VARCHAR(32) NOT NULL DEFAULT 'legacy'" ) if dialect == "postgresql": # Postgres needs the default, and we set it server-side so existing # rows back-fill automatically. return ( "ALTER TABLE maf_roundtable_seat_runs " "ADD COLUMN execution_stage VARCHAR(32) NOT NULL DEFAULT 'legacy'" ) # SQLite < 3.35 does not support adding a NOT NULL column with a default # via ALTER TABLE in the same way; but modern SQLite (3.35+, 2021+) does. # The default 'legacy' back-fills existing rows. return ( "ALTER TABLE maf_roundtable_seat_runs " "ADD COLUMN execution_stage VARCHAR(32) NOT NULL DEFAULT 'legacy'" ) def _ddl_add_result_snapshot(dialect: str) -> str: """Add the nullable ``result_snapshot`` JSON column (Phase-3 P0-1b). Stores a sanitized, replayable snapshot of the SeatResponse written atomically with the completed transition. Nullable so the ALTER is safe on all dialects (existing rows get NULL). """ if dialect == "mysql": return ( "ALTER TABLE maf_roundtable_seat_runs " "ADD COLUMN result_snapshot JSON NULL" ) if dialect == "postgresql": return ( "ALTER TABLE maf_roundtable_seat_runs " "ADD COLUMN result_snapshot JSON NULL" ) # SQLite — reached only when not rebuilding; use TEXT for JSON. return ( "ALTER TABLE maf_roundtable_seat_runs " "ADD COLUMN result_snapshot TEXT" ) def _ddl_drop_constraint(dialect: str, name: str) -> str: if dialect == "mysql": # MySQL uses DROP INDEX for both unique constraints and unique indexes. return f"ALTER TABLE maf_roundtable_seat_runs DROP INDEX {name}" # PostgreSQL / ANSI return f"ALTER TABLE maf_roundtable_seat_runs DROP CONSTRAINT {name}" def _ddl_drop_index(name: str) -> str: return f"DROP INDEX {name}" def _ddl_add_new_constraint(dialect: str) -> str: new = NEW_CONSTRAINT_NAME cols = "maf_run_id, agent_id, turn_index, execution_stage" if dialect == "mysql": return ( f"ALTER TABLE maf_roundtable_seat_runs " f"ADD UNIQUE INDEX {new} ({cols})" ) # PostgreSQL — table constraint return ( f"ALTER TABLE maf_roundtable_seat_runs " f"ADD CONSTRAINT {new} UNIQUE ({cols})" ) def _ddl_add_execution_stage_check(dialect: str) -> str: """Add the closed-enum CHECK constraint on execution_stage. Mirrors ``seat_runs.VALID_EXECUTION_STAGES`` so the DB itself rejects a typo (e.g. 'exprt') that would otherwise create a brand-new stage and bypass the four-column idempotency key. Keep the value list in sync with ``VALID_EXECUTION_STAGES``. """ allowed = ( "'legacy','planner','expert','format_repair'," "'cross_review','cross_review_repair','finalizer','group_chat'" ) if dialect == "mysql": # MySQL 8 accepts inline CHECK; older MySQL silently ignores it, but the # app-level validation in seat_runs._validate_stage is the real guard. return ( f"ALTER TABLE maf_roundtable_seat_runs " f"ADD CONSTRAINT ck_maf_seat_runs_execution_stage " f"CHECK (execution_stage IN ({allowed}))" ) # PostgreSQL / ANSI return ( f"ALTER TABLE maf_roundtable_seat_runs " f"ADD CONSTRAINT ck_maf_seat_runs_execution_stage " f"CHECK (execution_stage IN ({allowed}))" ) def _sqlite_rebuild_seat_runs( *, source_has_stage: bool = False, source_has_snapshot: bool = False, ) -> list[str]: """Rebuild the table in SQLite to swap the unique constraint. SQLite's ``ALTER TABLE`` cannot drop a constraint, so the portable recipe is: create a new table with the desired shape, copy data over, drop the old table, and rename. Indexes are recreated on the new table. ``source_has_stage`` / ``source_has_snapshot`` control the INSERT SELECT expressions so an intermediate table (already has stage and/or snapshot) does NOT lose data by being overwritten with defaults. """ allowed = ( "'legacy','planner','expert','format_repair'," "'cross_review','cross_review_repair','finalizer','group_chat'" ) source_stage = "execution_stage" if source_has_stage else "'legacy'" source_snapshot = "result_snapshot" if source_has_snapshot else "NULL" return [ f""" CREATE TABLE maf_roundtable_seat_runs__phase3 ( 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, execution_stage VARCHAR(32) NOT NULL DEFAULT 'legacy', 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, result_snapshot TEXT, 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, CONSTRAINT {NEW_CONSTRAINT_NAME} UNIQUE (maf_run_id, agent_id, turn_index, execution_stage), CONSTRAINT ck_maf_seat_runs_execution_stage CHECK (execution_stage IN ({allowed})) ) """, f""" INSERT INTO maf_roundtable_seat_runs__phase3 ( id, maf_run_id, agent_id, turn_index, execution_stage, deerflow_run_id, deerflow_thread_id, sandbox_id, workflow_checkpoint_id, status, result_snapshot, started_at, finished_at, error, created_at, updated_at ) SELECT id, maf_run_id, agent_id, turn_index, {source_stage}, deerflow_run_id, deerflow_thread_id, sandbox_id, workflow_checkpoint_id, status, {source_snapshot}, started_at, finished_at, error, created_at, updated_at FROM maf_roundtable_seat_runs """, "DROP TABLE maf_roundtable_seat_runs", "ALTER TABLE maf_roundtable_seat_runs__phase3 RENAME TO maf_roundtable_seat_runs", "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_deerflow_run_id ON maf_roundtable_seat_runs (deerflow_run_id)", "CREATE INDEX ix_maf_seat_runs_status ON maf_roundtable_seat_runs (status)", ] def main() -> int: parser = argparse.ArgumentParser(description="Phase 3 migration for MAF roundtable seat runs") 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())