alert-contract-expiring-hr/scripts/db/migrations.py

"""增量迁移：在 init_db 之外支持 idempotent ALTER/CREATE。

v1.1.1 新增：
- materials_content 表（材料包 PDF/Word 解析缓存）
- renewal_tasks.decision_source 列（AI 建议来源透明化）
- renewal_tasks.materials_summary 列（材料摘要缓存，便于离线渲染）
"""

from __future__ import annotations

import sqlite3


def _column_exists(conn: sqlite3.Connection, table: str, column: str) -> bool:
    rows = conn.execute(f"PRAGMA table_info({table})").fetchall()
    return any(r["name"] == column for r in rows)


def ensure_materials_content_schema(conn: sqlite3.Connection) -> None:
    """创建 materials_content 表（如不存在）。幂等。"""
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS materials_content (
            contract_id      TEXT NOT NULL,
            file_name        TEXT NOT NULL,
            file_path        TEXT NOT NULL,
            extracted_text   TEXT,
            parse_method     TEXT,
            parsed_at        TEXT NOT NULL,
            PRIMARY KEY (contract_id, file_name),
            FOREIGN KEY (contract_id) REFERENCES contracts(contract_id)
        );
        """
    )
    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_materials_contract "
        "ON materials_content(contract_id);"
    )


def ensure_renewal_tasks_columns(conn: sqlite3.Connection) -> None:
    """为 renewal_tasks 加 decision_source / materials_summary 列（如不存在）。"""
    if not _column_exists(conn, "renewal_tasks", "decision_source"):
        conn.execute(
            "ALTER TABLE renewal_tasks ADD COLUMN decision_source TEXT DEFAULT 'rules';"
        )
    if not _column_exists(conn, "renewal_tasks", "materials_summary"):
        conn.execute(
            "ALTER TABLE renewal_tasks ADD COLUMN materials_summary TEXT DEFAULT '';"
        )


def apply_all(conn: sqlite3.Connection) -> None:
    """应用全部 v1.1.1 迁移。"""
    ensure_materials_content_schema(conn)
    ensure_renewal_tasks_columns(conn)