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

"""CRUD 仓储：contracts / performance_records / alerts / renewal_tasks / dispatch_log。

所有方法接收一个外部传入的 sqlite3.Connection，调用方负责事务边界。
"""

from __future__ import annotations

import json
import sqlite3
from datetime import date, datetime
from typing import Any, Iterable


# ---------- contracts ----------

def list_active_contracts(conn: sqlite3.Connection) -> list[sqlite3.Row]:
    """列出 status='可续聘' 的合同。"""
    return conn.execute(
        "SELECT * FROM contracts WHERE status = '可续聘' ORDER BY contract_end ASC"
    ).fetchall()


def get_contract(conn: sqlite3.Connection, contract_id: str) -> sqlite3.Row | None:
    return conn.execute(
        "SELECT * FROM contracts WHERE contract_id = ?", (contract_id,)
    ).fetchone()


def update_contract_status(
    conn: sqlite3.Connection,
    contract_id: str,
    status: str,
    material_id: str | None = None,
) -> int:
    """更新合同状态。material_id 仅在 status='已续签' 时写入。"""
    renewed_at = datetime.now().isoformat(timespec="seconds") if status == "已续签" else None
    cur = conn.execute(
        """
        UPDATE contracts
        SET status = ?,
            last_material_id = COALESCE(?, last_material_id),
            renewed_at = COALESCE(?, renewed_at)
        WHERE contract_id = ?
        """,
        (status, material_id, renewed_at, contract_id),
    )
    return cur.rowcount


# ---------- performance ----------

def list_performance_by_contract(
    conn: sqlite3.Connection, contract_id: str
) -> list[sqlite3.Row]:
    """某合同的业绩记录，按年份升序。"""
    return conn.execute(
        "SELECT * FROM performance_records WHERE contract_id = ? ORDER BY year ASC",
        (contract_id,),
    ).fetchall()


# ---------- alerts ----------

def insert_alert(
    conn: sqlite3.Connection,
    alert_id: str,
    contract_id: str,
    level: str,
    days_to_expire: int,
    scan_id: str,
) -> None:
    conn.execute(
        """
        INSERT OR REPLACE INTO alerts
            (alert_id, contract_id, alert_level, days_to_expire, scan_id, generated_at, status)
        VALUES (?, ?, ?, ?, ?, ?, 'open')
        """,
        (
            alert_id,
            contract_id,
            level,
            days_to_expire,
            scan_id,
            datetime.now().isoformat(timespec="seconds"),
        ),
    )


def close_open_alerts_for_contract(conn: sqlite3.Connection, contract_id: str) -> int:
    """将某合同所有 open 状态的 alerts 关闭（关联 '已续签' 等终态）。"""
    cur = conn.execute(
        "UPDATE alerts SET status = 'closed' WHERE contract_id = ? AND status = 'open'",
        (contract_id,),
    )
    return cur.rowcount


def list_open_alerts_by_unit_month(
    conn: sqlite3.Connection, unit_id: str, month: str
) -> list[sqlite3.Row]:
    """按单位/月份列出关联合同的 open alerts（用于派发汇总）。"""
    sql = """
    SELECT a.alert_id, a.contract_id, a.alert_level, a.days_to_expire,
           c.contract_end, c.teacher_name, c.unit_name
    FROM alerts a
    JOIN contracts c ON c.contract_id = a.contract_id
    WHERE a.status = 'open'
      AND c.unit_id = ?
    """
    params: list[Any] = [unit_id]
    if month:
        sql += " AND substr(c.contract_end, 1, 7) = ?"
        params.append(month)
    sql += " ORDER BY c.contract_end ASC"
    return conn.execute(sql, params).fetchall()


# ---------- renewal_tasks ----------

def insert_task(
    conn: sqlite3.Connection,
    task_id: str,
    contract_id: str,
    alert_id: str,
    suggestion: str,
    reason: str,
    summary: str,
    checklist: list[str],
) -> None:
    conn.execute(
        """
        INSERT OR REPLACE INTO renewal_tasks
            (task_id, contract_id, alert_id, ai_suggestion, suggestion_reason,
             performance_summary, materials_checklist, generated_at, status)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending')
        """,
        (
            task_id,
            contract_id,
            alert_id,
            suggestion,
            reason,
            summary,
            json.dumps(checklist, ensure_ascii=False),
            datetime.now().isoformat(timespec="seconds"),
        ),
    )


def get_task(conn: sqlite3.Connection, task_id: str) -> sqlite3.Row | None:
    return conn.execute(
        """
        SELECT t.*, c.teacher_name, c.unit_name, c.position, c.contract_type,
               c.contract_start, c.contract_end, c.employee_id
        FROM renewal_tasks t
        JOIN contracts c ON c.contract_id = t.contract_id
        WHERE t.task_id = ?
        """,
        (task_id,),
    ).fetchone()


# ---------- dispatch_log ----------

def insert_dispatch_log(
    conn: sqlite3.Connection,
    log_id: str,
    unit_id: str,
    month: str,
    task_ids: list[str],
    payload_path: str,
    status: str = "dry_run",
) -> None:
    conn.execute(
        """
        INSERT INTO dispatch_log
            (log_id, unit_id, month, task_ids, dispatched_at, payload_path, status)
        VALUES (?, ?, ?, ?, ?, ?, ?)
        """,
        (
            log_id,
            unit_id,
            month,
            json.dumps(task_ids, ensure_ascii=False),
            datetime.now().isoformat(timespec="seconds"),
            payload_path,
            status,
        ),
    )


def list_dispatch_logs(conn: sqlite3.Connection) -> list[sqlite3.Row]:
    return conn.execute(
        "SELECT * FROM dispatch_log ORDER BY dispatched_at DESC"
    ).fetchall()