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

"""SQLite 连接 + DDL 定义 + init_db()。

5 张表：
- contracts: 合同台账
- performance_records: 近 3 年业绩记录
- alerts: 预警记录
- renewal_tasks: 续聘任务
- dispatch_log: 派发日志
"""

from __future__ import annotations

import sqlite3
from pathlib import Path
from typing import Iterable

from util.runtime_paths import ensure_dirs, get_db_path

# DDL：所有 IF NOT EXISTS，幂等可重入
DDL_STATEMENTS: tuple[str, ...] = (
    """
    CREATE TABLE IF NOT EXISTS contracts (
        contract_id        TEXT PRIMARY KEY,
        employee_id        TEXT NOT NULL,
        teacher_name       TEXT NOT NULL,
        unit_id            TEXT NOT NULL,
        unit_name          TEXT NOT NULL,
        position           TEXT NOT NULL,
        contract_type      TEXT NOT NULL,
        contract_start     TEXT NOT NULL,
        contract_end       TEXT NOT NULL,
        trial_period_end   TEXT NOT NULL,
        status             TEXT NOT NULL DEFAULT '可续聘',
        last_material_id   TEXT,
        renewed_at         TEXT
    );
    """,
    """
    CREATE TABLE IF NOT EXISTS performance_records (
        record_id              TEXT PRIMARY KEY,
        contract_id            TEXT NOT NULL,
        year                   INTEGER NOT NULL,
        performance_score      REAL NOT NULL,
        moral_teaching_score   REAL NOT NULL,
        research_score         REAL NOT NULL,
        teaching_score         REAL NOT NULL,
        FOREIGN KEY (contract_id) REFERENCES contracts(contract_id)
    );
    """,
    "CREATE INDEX IF NOT EXISTS idx_perf_contract ON performance_records(contract_id);",
    """
    CREATE TABLE IF NOT EXISTS alerts (
        alert_id        TEXT PRIMARY KEY,
        contract_id     TEXT NOT NULL,
        alert_level     TEXT NOT NULL,
        days_to_expire  INTEGER NOT NULL,
        scan_id         TEXT NOT NULL,
        generated_at    TEXT NOT NULL,
        status          TEXT NOT NULL DEFAULT 'open',
        FOREIGN KEY (contract_id) REFERENCES contracts(contract_id)
    );
    """,
    "CREATE INDEX IF NOT EXISTS idx_alerts_contract ON alerts(contract_id);",
    "CREATE INDEX IF NOT EXISTS idx_alerts_status ON alerts(status);",
    """
    CREATE TABLE IF NOT EXISTS renewal_tasks (
        task_id               TEXT PRIMARY KEY,
        contract_id           TEXT NOT NULL,
        alert_id              TEXT,
        ai_suggestion         TEXT NOT NULL,
        suggestion_reason     TEXT NOT NULL,
        performance_summary   TEXT NOT NULL,
        materials_checklist   TEXT NOT NULL,
        generated_at          TEXT NOT NULL,
        status                TEXT NOT NULL DEFAULT 'pending',
        FOREIGN KEY (contract_id) REFERENCES contracts(contract_id),
        FOREIGN KEY (alert_id)    REFERENCES alerts(alert_id)
    );
    """,
    "CREATE INDEX IF NOT EXISTS idx_tasks_contract ON renewal_tasks(contract_id);",
    """
    CREATE TABLE IF NOT EXISTS dispatch_log (
        log_id          TEXT PRIMARY KEY,
        unit_id         TEXT NOT NULL,
        month           TEXT NOT NULL,
        task_ids        TEXT NOT NULL,
        dispatched_at   TEXT NOT NULL,
        payload_path    TEXT NOT NULL,
        status          TEXT NOT NULL DEFAULT 'dry_run'
    );
    """,
    "CREATE INDEX IF NOT EXISTS idx_dispatch_unit_month ON dispatch_log(unit_id, month);",
)


def get_connection(db_path: Path | None = None) -> sqlite3.Connection:
    """获取 SQLite 连接。"""
    ensure_dirs()
    path = db_path or get_db_path()
    conn = sqlite3.connect(str(path))
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA foreign_keys = ON;")
    return conn


def init_db(db_path: Path | None = None) -> dict[str, int]:
    """建表 + 加载种子数据。返回表行数 dict。"""
    ensure_dirs()
    conn = get_connection(db_path)
    try:
        for ddl in DDL_STATEMENTS:
            conn.execute(ddl)
        conn.commit()

        # 加载种子（如果 contracts 表为空）
        from db.seed_data import load_seed_if_empty

        load_seed_if_empty(conn)

        # 统计
        counts = {
            "contracts": conn.execute("SELECT COUNT(*) AS n FROM contracts").fetchone()["n"],
            "performance_records": conn.execute("SELECT COUNT(*) AS n FROM performance_records").fetchone()["n"],
            "alerts": conn.execute("SELECT COUNT(*) AS n FROM alerts").fetchone()["n"],
            "renewal_tasks": conn.execute("SELECT COUNT(*) AS n FROM renewal_tasks").fetchone()["n"],
            "dispatch_log": conn.execute("SELECT COUNT(*) AS n FROM dispatch_log").fetchone()["n"],
        }
        return counts
    finally:
        conn.close()


def table_names() -> Iterable[str]:
    return ("contracts", "performance_records", "alerts", "renewal_tasks", "dispatch_log")