460 lines
20 KiB
Python
460 lines
20 KiB
Python
"""V2 migration for MAF roundtable conversation tables.
|
|
|
|
Creates the three user-visible conversation tables:
|
|
- maf_roundtable_conversations
|
|
- maf_roundtable_messages
|
|
- maf_roundtable_conversation_events
|
|
|
|
Idempotent: safe to run multiple times on SQLite / PostgreSQL / MySQL.
|
|
Does NOT modify or migrate existing MAF session/run/event data.
|
|
|
|
Usage:
|
|
PYTHONPATH=. uv run python scripts/migrate_maf_roundtable_conversation_v2.py --dry-run
|
|
PYTHONPATH=. uv run python scripts/migrate_maf_roundtable_conversation_v2.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_conversation_v2")
|
|
|
|
CONVERSATIONS_TABLE = "maf_roundtable_conversations"
|
|
MESSAGES_TABLE = "maf_roundtable_messages"
|
|
EVENTS_TABLE = "maf_roundtable_conversation_events"
|
|
|
|
|
|
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(dialect: str) -> str:
|
|
if dialect == "postgresql":
|
|
return "TIMESTAMP WITH TIME ZONE"
|
|
if dialect in {"mysql", "mariadb"}:
|
|
return "DATETIME(6)"
|
|
return "DATETIME"
|
|
|
|
|
|
def _json_sql(dialect: str) -> str:
|
|
if dialect == "postgresql":
|
|
return "JSONB"
|
|
if dialect in {"mysql", "mariadb"}:
|
|
return "JSON"
|
|
return "TEXT"
|
|
|
|
|
|
def _longtext_sql(dialect: str) -> str:
|
|
if dialect == "postgresql":
|
|
return "TEXT"
|
|
if dialect in {"mysql", "mariadb"}:
|
|
return "LONGTEXT"
|
|
return "TEXT"
|
|
|
|
|
|
def _conversations_ddl(dialect: str) -> str:
|
|
ts = _timestamp_sql(dialect)
|
|
js = _json_sql(dialect)
|
|
return f"""
|
|
CREATE TABLE {CONVERSATIONS_TABLE} (
|
|
id VARCHAR(64) NOT NULL PRIMARY KEY,
|
|
owner_user_id VARCHAR(64) NOT NULL,
|
|
title VARCHAR(255) NOT NULL DEFAULT '',
|
|
task_context_json {js},
|
|
safe_status VARCHAR(32) NOT NULL DEFAULT 'idle',
|
|
active_run_id VARCHAR(64),
|
|
internal_session_id VARCHAR(64),
|
|
last_stream_seq INTEGER NOT NULL DEFAULT 0,
|
|
created_at {ts} NOT NULL,
|
|
updated_at {ts} NOT NULL,
|
|
FOREIGN KEY (active_run_id) REFERENCES maf_roundtable_runs (id) ON DELETE SET NULL,
|
|
FOREIGN KEY (internal_session_id) REFERENCES maf_roundtable_sessions (id) ON DELETE SET NULL
|
|
)
|
|
"""
|
|
|
|
|
|
def _messages_ddl(dialect: str) -> str:
|
|
ts = _timestamp_sql(dialect)
|
|
js = _json_sql(dialect)
|
|
lt = _longtext_sql(dialect)
|
|
return f"""
|
|
CREATE TABLE {MESSAGES_TABLE} (
|
|
id VARCHAR(64) NOT NULL PRIMARY KEY,
|
|
conversation_id VARCHAR(64) NOT NULL,
|
|
seq INTEGER NOT NULL,
|
|
message_kind VARCHAR(32) NOT NULL DEFAULT 'user',
|
|
role VARCHAR(16) NOT NULL DEFAULT 'user',
|
|
agent_id VARCHAR(64),
|
|
run_id VARCHAR(64),
|
|
message_ref VARCHAR(191),
|
|
content {lt} NOT NULL DEFAULT '',
|
|
metadata_json {js},
|
|
created_at {ts} NOT NULL,
|
|
CONSTRAINT uq_maf_messages_conv_seq UNIQUE (conversation_id, seq),
|
|
FOREIGN KEY (conversation_id) REFERENCES {CONVERSATIONS_TABLE} (id) ON DELETE CASCADE
|
|
)
|
|
"""
|
|
|
|
|
|
def _events_ddl(dialect: str) -> str:
|
|
ts = _timestamp_sql(dialect)
|
|
js = _json_sql(dialect)
|
|
return f"""
|
|
CREATE TABLE {EVENTS_TABLE} (
|
|
id VARCHAR(64) NOT NULL PRIMARY KEY,
|
|
conversation_id VARCHAR(64) NOT NULL,
|
|
seq INTEGER NOT NULL,
|
|
run_id VARCHAR(64),
|
|
source_event_id VARCHAR(191),
|
|
event_type VARCHAR(64) NOT NULL,
|
|
data_json {js},
|
|
created_at {ts} NOT NULL,
|
|
CONSTRAINT uq_maf_conv_events_conv_seq UNIQUE (conversation_id, seq),
|
|
CONSTRAINT uq_maf_conv_events_conv_source UNIQUE (conversation_id, source_event_id),
|
|
FOREIGN KEY (conversation_id) REFERENCES {CONVERSATIONS_TABLE} (id) ON DELETE CASCADE
|
|
)
|
|
"""
|
|
|
|
|
|
def _indexes_ddl(dialect: str) -> list[str]:
|
|
stmts: list[str] = []
|
|
stmts.append(f"CREATE INDEX ix_maf_conv_owner_updated ON {CONVERSATIONS_TABLE} (owner_user_id, updated_at)")
|
|
stmts.append(f"CREATE INDEX ix_maf_conv_safe_status ON {CONVERSATIONS_TABLE} (safe_status)")
|
|
stmts.append(f"CREATE INDEX ix_maf_msg_conv_seq ON {MESSAGES_TABLE} (conversation_id, seq)")
|
|
stmts.append(f"CREATE INDEX ix_maf_msg_created ON {MESSAGES_TABLE} (created_at)")
|
|
stmts.append(f"CREATE INDEX ix_maf_msg_run ON {MESSAGES_TABLE} (run_id)")
|
|
stmts.append(f"CREATE INDEX ix_maf_msg_ref ON {MESSAGES_TABLE} (message_ref)")
|
|
stmts.append(f"CREATE INDEX ix_maf_cevt_conv_seq ON {EVENTS_TABLE} (conversation_id, seq)")
|
|
stmts.append(f"CREATE INDEX ix_maf_cevt_type ON {EVENTS_TABLE} (event_type)")
|
|
stmts.append(f"CREATE INDEX ix_maf_cevt_created ON {EVENTS_TABLE} (created_at)")
|
|
stmts.append(f"CREATE INDEX ix_maf_cevt_source ON {EVENTS_TABLE} (source_event_id)")
|
|
return stmts
|
|
|
|
|
|
def _events_table_has_source_event_id(inspector: Any) -> bool:
|
|
"""Check if the events table already has the source_event_id column."""
|
|
if EVENTS_TABLE not in inspector.get_table_names():
|
|
return False
|
|
cols = {c["name"] for c in inspector.get_columns(EVENTS_TABLE)}
|
|
return "source_event_id" in cols
|
|
|
|
|
|
def _events_table_has_source_unique(inspector: Any) -> bool:
|
|
"""Check if the events table already has the (conversation_id, source_event_id) unique constraint or index.
|
|
|
|
SQLite exposes ``CREATE UNIQUE INDEX`` via ``get_indexes()``, not
|
|
``get_unique_constraints()``. PostgreSQL and MySQL may use either path.
|
|
Both must be checked so the migration is truly idempotent on every
|
|
supported dialect.
|
|
"""
|
|
if EVENTS_TABLE not in inspector.get_table_names():
|
|
return False
|
|
target = {"conversation_id", "source_event_id"}
|
|
uq_cols = [set(u.get("column_names", [])) for u in inspector.get_unique_constraints(EVENTS_TABLE)]
|
|
if target in uq_cols:
|
|
return True
|
|
idx_cols = [
|
|
set(idx.get("column_names", []))
|
|
for idx in inspector.get_indexes(EVENTS_TABLE)
|
|
if idx.get("unique")
|
|
]
|
|
return target in idx_cols
|
|
|
|
|
|
def _alter_add_source_event_id_sql(dialect: str) -> str:
|
|
"""ALTER TABLE to add source_event_id column, portable across dialects."""
|
|
col_type = "VARCHAR(191)" if dialect in {"postgresql", "mysql", "mariadb"} else "VARCHAR(191)"
|
|
return f"ALTER TABLE {EVENTS_TABLE} ADD COLUMN source_event_id {col_type}"
|
|
|
|
|
|
def _create_source_unique_index_sql(dialect: str) -> str:
|
|
"""Create a unique index on (conversation_id, source_event_id).
|
|
|
|
Using CREATE UNIQUE INDEX instead of ALTER TABLE ADD CONSTRAINT because:
|
|
- SQLite only supports adding constraints via CREATE INDEX (no ALTER TABLE ADD CONSTRAINT)
|
|
- PostgreSQL and MySQL both accept CREATE UNIQUE INDEX
|
|
- Existing rows with NULL source_event_id are not deduped (SQL NULLs are distinct)
|
|
"""
|
|
return (
|
|
f"CREATE UNIQUE INDEX IF NOT EXISTS uq_maf_conv_events_conv_source "
|
|
f"ON {EVENTS_TABLE} (conversation_id, source_event_id)"
|
|
)
|
|
|
|
|
|
def validate_schema(inspector: Any) -> list[str]:
|
|
"""Return contract violations found by SQLAlchemy inspector."""
|
|
issues: list[str] = []
|
|
tables = inspector.get_table_names()
|
|
|
|
for tbl in (CONVERSATIONS_TABLE, MESSAGES_TABLE, EVENTS_TABLE):
|
|
if tbl not in tables:
|
|
issues.append(f"missing table: {tbl}")
|
|
|
|
if CONVERSATIONS_TABLE in tables:
|
|
cols = {c["name"] for c in inspector.get_columns(CONVERSATIONS_TABLE)}
|
|
for required in ("id", "owner_user_id", "title", "safe_status", "last_stream_seq", "created_at", "updated_at"):
|
|
if required not in cols:
|
|
issues.append(f"{CONVERSATIONS_TABLE}: missing column {required}")
|
|
if "last_message_seq" in cols:
|
|
issues.append(f"{CONVERSATIONS_TABLE}: deprecated column last_message_seq still present")
|
|
fks = inspector.get_foreign_keys(CONVERSATIONS_TABLE)
|
|
fk_columns = {
|
|
(tuple(fk.get("constrained_columns") or ()), fk.get("referred_table"))
|
|
for fk in fks
|
|
}
|
|
if (("active_run_id",), "maf_roundtable_runs") not in fk_columns:
|
|
issues.append(f"{CONVERSATIONS_TABLE}: missing FK(active_run_id -> maf_roundtable_runs)")
|
|
if (("internal_session_id",), "maf_roundtable_sessions") not in fk_columns:
|
|
issues.append(f"{CONVERSATIONS_TABLE}: missing FK(internal_session_id -> maf_roundtable_sessions)")
|
|
|
|
if MESSAGES_TABLE in tables:
|
|
cols = {c["name"] for c in inspector.get_columns(MESSAGES_TABLE)}
|
|
for required in ("id", "conversation_id", "seq", "message_kind", "role", "content", "created_at"):
|
|
if required not in cols:
|
|
issues.append(f"{MESSAGES_TABLE}: missing column {required}")
|
|
uqs = inspector.get_unique_constraints(MESSAGES_TABLE)
|
|
uq_cols = [set(u.get("column_names", [])) for u in uqs]
|
|
if {"conversation_id", "seq"} not in uq_cols:
|
|
issues.append(f"{MESSAGES_TABLE}: missing UNIQUE(conversation_id, seq)")
|
|
|
|
if EVENTS_TABLE in tables:
|
|
cols = {c["name"] for c in inspector.get_columns(EVENTS_TABLE)}
|
|
for required in ("id", "conversation_id", "seq", "event_type", "created_at", "source_event_id"):
|
|
if required not in cols:
|
|
issues.append(f"{EVENTS_TABLE}: missing column {required}")
|
|
uqs = inspector.get_unique_constraints(EVENTS_TABLE)
|
|
uq_cols = [set(u.get("column_names", [])) for u in uqs]
|
|
# Also check unique indexes — SQLite exposes CREATE UNIQUE INDEX via
|
|
# get_indexes(), not get_unique_constraints().
|
|
idx_cols = [
|
|
set(idx.get("column_names", []))
|
|
for idx in inspector.get_indexes(EVENTS_TABLE)
|
|
if idx.get("unique")
|
|
]
|
|
all_uq = uq_cols + idx_cols
|
|
if {"conversation_id", "seq"} not in all_uq:
|
|
issues.append(f"{EVENTS_TABLE}: missing UNIQUE(conversation_id, seq)")
|
|
if {"conversation_id", "source_event_id"} not in all_uq:
|
|
issues.append(f"{EVENTS_TABLE}: missing UNIQUE(conversation_id, source_event_id)")
|
|
|
|
return issues
|
|
|
|
|
|
def _missing_conversation_link_foreign_keys(inspector: Any) -> list[str]:
|
|
"""Return missing private-link FKs on an already-created conversation table."""
|
|
if CONVERSATIONS_TABLE not in inspector.get_table_names():
|
|
return []
|
|
foreign_keys = {
|
|
(tuple(fk.get("constrained_columns") or ()), fk.get("referred_table"))
|
|
for fk in inspector.get_foreign_keys(CONVERSATIONS_TABLE)
|
|
}
|
|
missing: list[str] = []
|
|
if (("active_run_id",), "maf_roundtable_runs") not in foreign_keys:
|
|
missing.append("active_run_id -> maf_roundtable_runs")
|
|
if (("internal_session_id",), "maf_roundtable_sessions") not in foreign_keys:
|
|
missing.append("internal_session_id -> maf_roundtable_sessions")
|
|
return missing
|
|
|
|
|
|
async def run_migration(database_url: str, *, dry_run: bool = False) -> None:
|
|
from sqlalchemy import inspect, text
|
|
from sqlalchemy.ext.asyncio import create_async_engine
|
|
from sqlalchemy.pool import NullPool
|
|
|
|
# A migration is a one-shot operation. NullPool makes every early
|
|
# failure close its SQLite connection as the transaction context exits,
|
|
# rather than retaining it until the process exits.
|
|
engine = create_async_engine(database_url, poolclass=NullPool)
|
|
async with engine.begin() as conn:
|
|
dialect = _dialect_name(conn)
|
|
logger.info("Dialect: %s", dialect)
|
|
|
|
def _get_table_names(sync_conn: Any) -> list[str]:
|
|
return inspect(sync_conn).get_table_names()
|
|
|
|
existing = set(await conn.run_sync(_get_table_names))
|
|
|
|
def _missing_links(sync_conn: Any) -> list[str]:
|
|
return _missing_conversation_link_foreign_keys(inspect(sync_conn))
|
|
|
|
missing_links = await conn.run_sync(_missing_links)
|
|
if missing_links:
|
|
if dialect in {"sqlite"}:
|
|
# SQLite cannot add a foreign-key constraint in place. These
|
|
# V2 tables have not been released as a data source yet, so an
|
|
# empty partial install can be rebuilt safely. Never delete
|
|
# a non-empty conversation history merely to improve a schema.
|
|
tables_to_rebuild = [EVENTS_TABLE, MESSAGES_TABLE, CONVERSATIONS_TABLE]
|
|
non_empty: list[str] = []
|
|
for table in tables_to_rebuild:
|
|
if table not in existing:
|
|
continue
|
|
count = int((await conn.execute(text(f"SELECT COUNT(*) FROM {table}"))).scalar_one())
|
|
if count:
|
|
non_empty.append(table)
|
|
if non_empty:
|
|
# ``run_migration`` normally disposes below the transaction
|
|
# block. This explicit failure is intentionally raised
|
|
# earlier, so release the aiosqlite pool before bubbling it
|
|
# to callers/tests.
|
|
await engine.dispose()
|
|
raise RuntimeError(
|
|
"MAF Conversation V2 schema is missing foreign keys "
|
|
f"({', '.join(missing_links)}) but has data in "
|
|
f"{', '.join(non_empty)}. Back up and migrate it explicitly; "
|
|
"this script refuses to delete conversation history."
|
|
)
|
|
logger.warning(
|
|
"Rebuilding empty partial V2 tables to add missing foreign keys: %s",
|
|
", ".join(missing_links),
|
|
)
|
|
if not dry_run:
|
|
for table in tables_to_rebuild:
|
|
if table in existing:
|
|
await conn.execute(text(f"DROP TABLE {table}"))
|
|
existing.difference_update(tables_to_rebuild)
|
|
else:
|
|
# PostgreSQL and MySQL/MariaDB can add the missing constraints
|
|
# in place. The database validates existing rows and fails
|
|
# closed if a dangling private link exists.
|
|
if "active_run_id -> maf_roundtable_runs" in missing_links:
|
|
sql = (
|
|
f"ALTER TABLE {CONVERSATIONS_TABLE} "
|
|
"ADD CONSTRAINT fk_maf_conv_active_run "
|
|
"FOREIGN KEY (active_run_id) REFERENCES maf_roundtable_runs (id) "
|
|
"ON DELETE SET NULL"
|
|
)
|
|
if dry_run:
|
|
logger.info("[DRY RUN] Would execute: %s", sql)
|
|
else:
|
|
await conn.execute(text(sql))
|
|
if "internal_session_id -> maf_roundtable_sessions" in missing_links:
|
|
sql = (
|
|
f"ALTER TABLE {CONVERSATIONS_TABLE} "
|
|
"ADD CONSTRAINT fk_maf_conv_internal_session "
|
|
"FOREIGN KEY (internal_session_id) REFERENCES maf_roundtable_sessions (id) "
|
|
"ON DELETE SET NULL"
|
|
)
|
|
if dry_run:
|
|
logger.info("[DRY RUN] Would execute: %s", sql)
|
|
else:
|
|
await conn.execute(text(sql))
|
|
|
|
ddls: list[str] = []
|
|
if CONVERSATIONS_TABLE not in existing:
|
|
ddls.append(_conversations_ddl(dialect))
|
|
logger.info("Will create: %s", CONVERSATIONS_TABLE)
|
|
else:
|
|
logger.info("Already exists: %s", CONVERSATIONS_TABLE)
|
|
|
|
if MESSAGES_TABLE not in existing:
|
|
ddls.append(_messages_ddl(dialect))
|
|
logger.info("Will create: %s", MESSAGES_TABLE)
|
|
else:
|
|
logger.info("Already exists: %s", MESSAGES_TABLE)
|
|
|
|
if EVENTS_TABLE not in existing:
|
|
ddls.append(_events_ddl(dialect))
|
|
logger.info("Will create: %s", EVENTS_TABLE)
|
|
else:
|
|
logger.info("Already exists: %s", EVENTS_TABLE)
|
|
|
|
# ── In-place upgrade for existing events table ──
|
|
# Old V2 deployments created the events table WITHOUT source_event_id.
|
|
# Detect and add the column + unique index so dedup works.
|
|
if EVENTS_TABLE in existing:
|
|
def _check_source_col(sync_conn: Any) -> bool:
|
|
return _events_table_has_source_event_id(inspect(sync_conn))
|
|
|
|
has_source_col = await conn.run_sync(_check_source_col)
|
|
if not has_source_col:
|
|
alter_sql = _alter_add_source_event_id_sql(dialect)
|
|
logger.info("Upgrading %s: adding source_event_id column", EVENTS_TABLE)
|
|
if dry_run:
|
|
logger.info("[DRY RUN] Would execute: %s", alter_sql)
|
|
else:
|
|
await conn.execute(text(alter_sql))
|
|
|
|
def _check_source_unique(sync_conn: Any) -> bool:
|
|
return _events_table_has_source_unique(inspect(sync_conn))
|
|
|
|
has_source_uq = await conn.run_sync(_check_source_unique)
|
|
if not has_source_uq:
|
|
uq_sql = _create_source_unique_index_sql(dialect)
|
|
logger.info("Upgrading %s: adding unique index (conversation_id, source_event_id)", EVENTS_TABLE)
|
|
if dry_run:
|
|
logger.info("[DRY RUN] Would execute: %s", uq_sql)
|
|
else:
|
|
try:
|
|
await conn.execute(text(uq_sql))
|
|
except Exception as exc:
|
|
if "already exists" in str(exc).lower():
|
|
pass
|
|
else:
|
|
raise
|
|
|
|
for idx_sql in _indexes_ddl(dialect):
|
|
ddls.append(idx_sql)
|
|
|
|
if dry_run:
|
|
logger.info("[DRY RUN] Would execute %d statements:", len(ddls))
|
|
for ddl in ddls:
|
|
logger.info(" %s", " ".join(ddl.split())[:120])
|
|
else:
|
|
for ddl in ddls:
|
|
try:
|
|
await conn.execute(text(ddl))
|
|
except Exception as exc:
|
|
if "already exists" in str(exc).lower() or "duplicate" in str(exc).lower():
|
|
continue
|
|
raise
|
|
logger.info("Executed %d statements successfully.", len(ddls))
|
|
|
|
def _validate(sync_conn: Any) -> list[str]:
|
|
return validate_schema(inspect(sync_conn))
|
|
|
|
issues = await conn.run_sync(_validate)
|
|
if issues:
|
|
logger.error("Schema validation FAILED:")
|
|
for issue in issues:
|
|
logger.error(" - %s", issue)
|
|
sys.exit(1)
|
|
else:
|
|
logger.info("Schema validation passed.")
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="MAF roundtable conversation V2 migration")
|
|
parser.add_argument("--database-url", default=None, help="Async database URL (default: from config)")
|
|
parser.add_argument("--dry-run", action="store_true", help="Print DDL without executing")
|
|
args = parser.parse_args()
|
|
|
|
database_url = args.database_url
|
|
if database_url is None:
|
|
from pathlib import Path
|
|
|
|
import yaml
|
|
|
|
config_path = Path(__file__).resolve().parent.parent / "config.yaml"
|
|
if config_path.exists():
|
|
with open(config_path) as f:
|
|
cfg = yaml.safe_load(f) or {}
|
|
db_cfg = cfg.get("database", {})
|
|
database_url = db_cfg.get("url", "sqlite+aiosqlite:///./data/deerflow.db")
|
|
else:
|
|
database_url = "sqlite+aiosqlite:///./data/deerflow.db"
|
|
|
|
logger.info("Database: %s", database_url.split("?")[0])
|
|
asyncio.run(run_migration(database_url, dry_run=args.dry_run))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|