attachment_classifier/db.py

# -*- coding: utf-8 -*-
"""SQL Server 数据访问层，通过 pyodbc 按总排号批量查询新参数字段。"""
from __future__ import annotations

import logging
from typing import Any

import pyodbc

logger = logging.getLogger(__name__)


class DatabaseError(Exception):
    """数据库连接或查询失败。"""


def _build_conn_str(db_cfg: dict[str, Any]) -> str:
    parts = [
        f"DRIVER={{{db_cfg['driver']}}};",
        f"SERVER={db_cfg['server']},{db_cfg['port']};",
        f"DATABASE={db_cfg['database']};",
        f"UID={db_cfg['username']};",
        f"PWD={db_cfg['password']};",
        f"Connection Timeout={db_cfg['connect_timeout']};",
    ]
    # 数据库服务器使用自签名/不受信任证书时，跳过证书链校验（连接仍保持加密）。
    # ODBC Driver 17/18 默认 Encrypt=Yes，遇到自签证书会报"不受信任的颁发机构"，
    # 加 TrustServerCertificate=Yes 即可信任该证书。设为 false 时不影响原有行为。
    if db_cfg.get("trust_server_certificate", False):
        parts.append("TrustServerCertificate=Yes;")
    return "".join(parts)


def fetch_params_by_ids(
    db_cfg: dict[str, Any], zong_pai_hao_list: list[str]
) -> dict[str, str | None]:
    """按总排号批量查询新参数字段。

    返回 dict：{总排号: 新参数文本}。数据库中不存在的总排号，其值为 None
    （而不是直接从结果里省略该 key），方便调用方区分"没查到"和"查到但内容为空"。

    一个总排号只对应一条记录（业务已确认为一对一关系）；如果实际数据出现
    重复总排号，取查询结果的第一条并记录一条 WARNING 日志，不中断整体流程。
    """
    if not zong_pai_hao_list:
        return {}

    result: dict[str, str | None] = {zph: None for zph in zong_pai_hao_list}

    schema = db_cfg["schema"]
    table = db_cfg["table"]
    id_col = db_cfg["id_column"]
    param_col = db_cfg["param_column"]

    # 表名必须带 schema 前缀（如 [dbo].[表名]），只写表名在 schema 不是默认dbo时
    # 会查到错误的表，甚至直接报"找不到对象"。schema 和表名分别加中括号转义，
    # 不能写成 [schema.table]，那样会被当成一个整体标识符解析。
    qualified_table = f"[{schema}].[{table}]"

    # 用参数化查询防止总排号里混入特殊字符导致 SQL 注入或语法错误
    placeholders = ",".join("?" for _ in zong_pai_hao_list)
    sql = f"SELECT [{id_col}], [{param_col}] FROM {qualified_table} WHERE [{id_col}] IN ({placeholders})"

    conn_str = _build_conn_str(db_cfg)
    try:
        with pyodbc.connect(conn_str, timeout=db_cfg["connect_timeout"]) as conn:
            cursor = conn.cursor()
            # 查询超时设在 Connection 上（pyodbc 的 timeout 是 Connection 属性，
            # Cursor 没有该属性，设 cursor.timeout 会报 AttributeError）。
            conn.timeout = db_cfg["query_timeout"]
            cursor.execute(sql, zong_pai_hao_list)
            seen = set()
            for row in cursor.fetchall():
                zph, param = row[0], row[1]
                if zph in seen:
                    logger.warning("总排号 %s 存在重复记录，已取第一条", zph)
                    continue
                seen.add(zph)
                result[zph] = param
    except pyodbc.Error as e:
        raise DatabaseError(f"数据库查询失败: {e}") from e

    return result


def fetch_param_by_id(db_cfg: dict[str, Any], zong_pai_hao: str) -> str | None:
    """单个总排号查询的便捷封装。"""
    return fetch_params_by_ids(db_cfg, [zong_pai_hao]).get(zong_pai_hao)


# ---------------------------------------------------------------------------
# 附件分类落库（目标表 Common.Attachment，与源表 productionContractData 分属不同 schema）
# ---------------------------------------------------------------------------
# 落库约定（与项目记忆 MEMORY.md 中写入约定一致，请勿在此处用字面量以外的值）：
#   - 无附件（has_attachment=false）：写哨兵行 (NS, '无附件', '无')
#   - 无小类的行（coarse 真实附件 / 哨兵行）：MinorCategory 统一填 '无'
ATTACHMENT_SCHEMA = "Common"
ATTACHMENT_TABLE = "Attachment"

SENTINEL_MAJOR = "无附件"
SENTINEL_MINOR = "无"
NO_MINOR = "无"


def fetch_all_ids(db_cfg: dict[str, Any], limit: int | None = None) -> list[str]:
    """拉取源表全部总排号（去重、排除 NULL），用于全量分类写入。

    limit 用于测试时只取前 N 个（按总排号排序）。返回的总排号列表可直接
    喂给 classify_batch。
    """
    schema = db_cfg["schema"]
    table = db_cfg["table"]
    id_col = db_cfg["id_column"]
    qualified = f"[{schema}].[{table}]"
    sql = f"SELECT DISTINCT [{id_col}] FROM {qualified} WHERE [{id_col}] IS NOT NULL"
    params: list[Any] = []
    if limit is not None:
        sql += f" ORDER BY [{id_col}] OFFSET 0 ROWS FETCH NEXT ? ROWS ONLY"
        params.append(int(limit))
    conn_str = _build_conn_str(db_cfg)
    try:
        with pyodbc.connect(conn_str, timeout=db_cfg["connect_timeout"]) as conn:
            conn.timeout = db_cfg["query_timeout"]
            cursor = conn.cursor()
            cursor.execute(sql, params)
            return [row[0] for row in cursor.fetchall()]
    except pyodbc.Error as e:
        raise DatabaseError(f"查询总排号失败: {e}") from e


def upsert_attachments(
    db_cfg: dict[str, Any],
    rows: list[tuple[str, str, str]],
) -> tuple[int, int]:
    """幂等写入 Common.Attachment。

    rows: 本次要写入的 (NS, MajorCategory, MinorCategory) 列表。
    对出现的每个 NS 先 DELETE 其旧行，再批量 INSERT——保证重跑总是反映
    最新分类结果，不会因唯一索引 (NS, 大类, 小类) 冲突而失败。

    返回 (deleted_rows, inserted_rows) 计数。
    """
    if not rows:
        return (0, 0)

    qualified = f"[{ATTACHMENT_SCHEMA}].[{ATTACHMENT_TABLE}]"
    distinct_ns = sorted({r[0] for r in rows})

    conn_str = _build_conn_str(db_cfg)
    try:
        with pyodbc.connect(conn_str, timeout=db_cfg["connect_timeout"]) as conn:
            conn.timeout = db_cfg["query_timeout"]
            cursor = conn.cursor()
            # 1) 删除本批所有 NS 的旧行
            del_ph = ",".join("?" for _ in distinct_ns)
            cursor.execute(
                f"DELETE FROM {qualified} WHERE [NS] IN ({del_ph})", distinct_ns
            )
            deleted = cursor.rowcount
            # 2) 插入新行
            cursor.executemany(
                f"INSERT INTO {qualified} ([NS], [MajorCategory], [MinorCategory]) "
                f"VALUES (?, ?, ?)",
                rows,
            )
            inserted = len(rows)
            conn.commit()
            logger.info(
                "upsert_attachments: 删除 %d 行, 插入 %d 行, 涉及 %d 个 NS",
                deleted, inserted, len(distinct_ns),
            )
            return (deleted, inserted)
    except pyodbc.Error as e:
        raise DatabaseError(f"写入附件表失败: {e}") from e