baodan/api/insurance/db/migrate_022.py

"""迁移 022: 工作区与统一任务表。

按 docs/0728修复文件.md 阶段 1 实施：
- 扩展 insurance_ppt_sessions：新增工作区字段
- 扩展 poster_records：新增工作区字段
- 新建 insurance_generation_tasks：统一任务执行表

兼容 PostgreSQL 和 MySQL。
"""
import logging
from sqlalchemy import text

logger = logging.getLogger(__name__)


def _column_exists(db, table: str, column: str) -> bool:
    """检查列是否存在（兼容 PostgreSQL 和 MySQL）。"""
    try:
        result = db.session.execute(text(
            "SELECT COUNT(*) FROM information_schema.columns "
            "WHERE table_name = :table AND column_name = :column"
        ), {"table": table, "column": column})
        return result.scalar() > 0
    except Exception:
        return False


def _table_exists(db, table_name: str) -> bool:
    """检查表是否存在（兼容 PostgreSQL 和 MySQL）。"""
    try:
        result = db.session.execute(text(
            "SELECT COUNT(*) FROM information_schema.tables "
            "WHERE table_name = :table"
        ), {"table": table_name})
        return result.scalar() > 0
    except Exception:
        return False


def _index_exists(db, index_name: str) -> bool:
    """检查索引是否存在（兼容 PostgreSQL 和 MySQL）。"""
    try:
        result = db.session.execute(text(
            "SELECT COUNT(*) FROM pg_indexes WHERE indexname = :name"
        ), {"name": index_name})
        return result.scalar() > 0
    except Exception:
        return False


def migrate():
    """执行迁移。"""
    from insurance.db.compat import db

    # ─── 扩展 PptSession ─────────────────────────────────────
    ppt_columns = [
        ("title", "VARCHAR(200) DEFAULT ''"),
        ("workflow_step", "VARCHAR(20) DEFAULT 'upload'"),
        ("draft_options_json", "TEXT"),
        ("draft_revision", "INT DEFAULT 1"),
        ("generated_revision", "INT DEFAULT 0"),
        ("latest_task_id", "VARCHAR(36)"),
        ("latest_output_path", "VARCHAR(500)"),
        ("archived_at", "TIMESTAMP"),
    ]

    for col_name, col_def in ppt_columns:
        if not _column_exists(db, "insurance_ppt_sessions", col_name):
            db.session.execute(text(
                f"ALTER TABLE insurance_ppt_sessions ADD COLUMN {col_name} {col_def}"
            ))
            logger.info(f"[migrate_022] 已为 insurance_ppt_sessions 添加 {col_name} 列")

    # ─── 扩展 PosterRecord ───────────────────────────────────
    poster_columns = [
        ("title", "VARCHAR(200) DEFAULT ''"),
        ("workflow_step", "VARCHAR(20) DEFAULT 'product'"),
        ("draft_status", "VARCHAR(20) DEFAULT 'active'"),
        ("draft_revision", "INT DEFAULT 1"),
        ("generated_revision", "INT DEFAULT 0"),
        ("latest_task_id", "VARCHAR(36)"),
        ("archived_at", "TIMESTAMP"),
    ]

    for col_name, col_def in poster_columns:
        if not _column_exists(db, "poster_records", col_name):
            db.session.execute(text(
                f"ALTER TABLE poster_records ADD COLUMN {col_name} {col_def}"
            ))
            logger.info(f"[migrate_022] 已为 poster_records 添加 {col_name} 列")

    # ─── 新建 insurance_generation_tasks ─────────────────────
    if not _table_exists(db, "insurance_generation_tasks"):
        db.session.execute(text("""
            CREATE TABLE insurance_generation_tasks (
                id VARCHAR(36) PRIMARY KEY,
                user_id VARCHAR(64) NOT NULL,
                artifact_type VARCHAR(10) NOT NULL,
                operation VARCHAR(10) NOT NULL,
                workspace_id VARCHAR(36) NOT NULL,
                title_snapshot VARCHAR(200) DEFAULT '',
                status VARCHAR(20) NOT NULL DEFAULT 'queued',
                stage VARCHAR(30) DEFAULT '',
                progress INT DEFAULT 0,
                message VARCHAR(500) DEFAULT '',
                error_code VARCHAR(50) DEFAULT '',
                error_message TEXT,
                input_revision INT DEFAULT 1,
                input_snapshot_json TEXT,
                output_json TEXT,
                idempotency_key VARCHAR(100),
                celery_task_id VARCHAR(200),
                attempt_count INT DEFAULT 0,
                heartbeat_at TIMESTAMP,
                started_at TIMESTAMP,
                finished_at TIMESTAMP,
                viewed_at TIMESTAMP,
                dock_hidden_at TIMESTAMP,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """))
        logger.info("[migrate_022] 已创建 insurance_generation_tasks 表")

        # 创建索引（PostgreSQL 兼容）
        indexes = [
            ("idx_gen_tasks_user_status", "user_id, status"),
            ("idx_gen_tasks_user_artifact", "user_id, artifact_type"),
            ("idx_gen_tasks_workspace", "workspace_id"),
            ("idx_gen_tasks_idempotency", "idempotency_key"),
        ]
        for idx_name, idx_cols in indexes:
            if not _index_exists(db, idx_name):
                db.session.execute(text(
                    f"CREATE INDEX {idx_name} ON insurance_generation_tasks ({idx_cols})"
                ))
                logger.info(f"[migrate_022] 已创建索引 {idx_name}")

    db.session.commit()
    logger.info("[migrate_022] 工作区与任务表迁移完成")