403 lines
16 KiB
Python
403 lines
16 KiB
Python
"""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())
|