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

"""Phase 5 migration for MAF external-write target leases.

Creates the additive ``maf_roundtable_action_leases`` table used to serialize
approved writes to the same server-derived external business target.  Existing
action rows are intentionally not backfilled: old rows have no trusted target
key and retain the legacy action-level at-most-once behavior.  New writes
receive an opaque target key from ``ActionCoordinator``.

Run this before deploying code that enables MAF external writes:

    PYTHONPATH=. uv run python scripts/migrate_maf_roundtable_phase5.py --dry-run
    PYTHONPATH=. uv run python scripts/migrate_maf_roundtable_phase5.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_phase5")

TABLE = "maf_roundtable_action_leases"
ACTION_INDEX = "ix_maf_action_leases_action_id"
EXPIRY_INDEX = "ix_maf_action_leases_lease_expires_at"


def _dialect_name(conn_or_dialect: Any) -> str:
    return str(getattr(getattr(conn_or_dialect, "dialect", conn_or_dialect), "name", "sqlite")).lower()


def _timestamp_sql(conn_or_dialect: Any) -> str:
    """Return the timestamp type accepted by each supported database."""

    dialect = _dialect_name(conn_or_dialect)
    if dialect == "postgresql":
        return "TIMESTAMP WITH TIME ZONE"
    if dialect in {"mysql", "mariadb"}:
        return "DATETIME(6)"
    return "DATETIME"


def _create_lease_table_sql(conn_or_dialect: Any) -> str:
    timestamp_type = _timestamp_sql(conn_or_dialect)
    return f"""
        CREATE TABLE {TABLE} (
            target_key VARCHAR(255) NOT NULL PRIMARY KEY,
            action_id VARCHAR(64) NOT NULL,
            lease_expires_at {timestamp_type} NOT NULL,
            created_at {timestamp_type} NOT NULL,
            updated_at {timestamp_type} NOT NULL,
            CONSTRAINT uq_maf_action_leases_action UNIQUE (action_id),
            FOREIGN KEY (action_id) REFERENCES maf_roundtable_actions (id) ON DELETE CASCADE
        )
    """


def _add_missing_columns(conn: Any, *, inspector: Any, execute: Any) -> None:
    """Complete a partially created table without rewriting existing data.

    The migration is additive. A missing primary key or action foreign-key
    column cannot be safely added to a populated table, so fail loudly rather
    than creating a lease table that the arbiter could misinterpret.
    """

    columns = {column["name"] for column in inspector.get_columns(TABLE)}
    required = {"target_key", "action_id", "lease_expires_at", "created_at", "updated_at"}
    missing = required - columns
    if not missing:
        return
    if {"target_key", "action_id"} & missing:
        raise RuntimeError(
            f"{TABLE} is incomplete ({sorted(missing)}); recreate it with the deployment migration"
        )
    timestamp_type = _timestamp_sql(conn)
    for column in ("lease_expires_at", "created_at", "updated_at"):
        if column in missing:
            execute(f"ALTER TABLE {TABLE} ADD COLUMN {column} {timestamp_type}")


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

    await init_engine_from_config(get_app_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"]
    async with engine.begin() as conn:
        await conn.run_sync(_migrate_sync, dry_run=dry_run)
    await close_engine()
    return 0


def _migrate_sync(conn: Any, *, dry_run: bool) -> None:
    from sqlalchemy import inspect, text

    inspector = inspect(conn)

    def execute(sql: str) -> None:
        if dry_run:
            logger.info("[dry-run] would execute: %s", " ".join(sql.split()))
            return
        conn.execute(text(sql))

    if TABLE not in inspector.get_table_names():
        logger.info("Creating %s", TABLE)
        execute(_create_lease_table_sql(conn))
        if dry_run:
            execute(f"CREATE INDEX {ACTION_INDEX} ON {TABLE} (action_id)")
            execute(f"CREATE INDEX {EXPIRY_INDEX} ON {TABLE} (lease_expires_at)")
            return
        inspector = inspect(conn)
    else:
        logger.info("Table %s already exists", TABLE)
        _add_missing_columns(conn, inspector=inspector, execute=execute)
        if not dry_run:
            inspector = inspect(conn)

    indexes = {index.get("name") for index in inspector.get_indexes(TABLE)}
    if ACTION_INDEX not in indexes:
        execute(f"CREATE INDEX {ACTION_INDEX} ON {TABLE} (action_id)")
    if EXPIRY_INDEX not in indexes:
        execute(f"CREATE INDEX {EXPIRY_INDEX} ON {TABLE} (lease_expires_at)")


def main() -> int:
    parser = argparse.ArgumentParser(description="Create MAF external-write target lease table")
    parser.add_argument("--dry-run", action="store_true", help="Preview only; do not write")
    return asyncio.run(_migrate(parser.parse_args().dry_run))


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