attachment_classifier/db.py

# -*- coding: utf-8 -*-
"""数据访问层：基于 SQLAlchemy 的数据库无关实现（支持 MSSQL / PostgreSQL）。

原 pyodbc 专用实现已重构为 SQLAlchemy Core + ORM：
  - 源表（列名/表名由配置驱动、含中文/动态）用 Core 的 table()/column() + quote=True
    动态构造，由 SQLAlchemy 按方言生成正确标识符引用（PG 用 "名"，MSSQL 用 [名]）。
  - 目标附件表 Common.Attachment 用 orm.Attachment（声明式 ORM 模型）读写/建表。

对外暴露的 4 个函数签名与返回结构与旧版完全一致，classifier.py /
write_attachments.py 无需改动调用方式（仅落库表结构/连接信息随配置变化）。
"""
from __future__ import annotations

import logging
from typing import Any

from sqlalchemy import Column, Integer, MetaData, String, Table, delete, insert, select
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.orm import Session

from orm import (
    Attachment,
    DatabaseError,
    SENTINEL_MAJOR,
    SENTINEL_MINOR,
    NO_MINOR,
    ATTACHMENT_SCHEMA,
    ATTACHMENT_TABLE,
    get_engine,
    init_schema,
)

logger = logging.getLogger(__name__)

# 兼容旧导入：把常量与异常从 orm 重新导出
__all__ = [
    "DatabaseError",
    "ATTACHMENT_SCHEMA",
    "ATTACHMENT_TABLE",
    "SENTINEL_MAJOR",
    "SENTINEL_MINOR",
    "NO_MINOR",
    "fetch_params_by_ids",
    "fetch_param_by_id",
    "fetch_all_ids",
    "fetch_existing_attachment_sns",
    "upsert_attachments",
    "init_schema",
]


# SQL Server 单条语句最多 2100 个绑定参数，WHERE col IN (?, ?, ...) 每项占 1 个，
# 超过会报 "COUNT 字段不正确或语法错误"。这里把超长 IN 列表按 2000 一批分片查询
# （2000 < 2100，留安全余量；对 PostgreSQL 的 65535 上限也完全够用）。
_PARAM_CHUNK = 2000


def _chunked(seq: list, size: int = _PARAM_CHUNK):
    """把列表切成 <= size 的连续片段，用于分批 IN 查询。"""
    for i in range(0, len(seq), size):
        yield seq[i : i + size]


def _build_source_table(db_cfg: dict[str, Any]):
    """按配置动态构造源表 Core 结构（带 quote，保证跨库标识符正确引用）。"""
    schema = db_cfg["schema"]
    tbl = db_cfg["table"]
    param_col = db_cfg["param_column"]
    sn_col = db_cfg["id_column"]
    id_field = db_cfg.get("id_field")

    cols = [
        Column(sn_col, String, quote=True),
        Column(param_col, String, quote=True),
    ]
    if id_field:
        cols.append(Column(id_field, Integer, quote=True))
    meta = MetaData()
    return Table(tbl, meta, *cols, schema=schema, quote=True, quote_schema=True)


def fetch_params_by_ids(
    db_cfg: dict[str, Any], id_keys: list[tuple[str, str]]
) -> dict[str, dict[str, Any | None]]:
    """按 (标识符, 键类型) 批量查询新参数字段。

    id_keys: list[(identifier, key)]，key 取值：
        "sn" -> 按总排号列（配置 id_column）查询，标识符即总排号；
        "id" -> 按数据库真实 ID 列（配置 id_field）查询，并回取对应的总排号。

    返回 dict：{identifier: {"param": 新参数文本|None, "sn": 总排号|None}}。
    标识符作为 key 原样保留（便于回查）；未查到的标识符其 param/sn 为 None。
    sn 为对应的总排号（键类型为 sn 时即标识符本身，键类型为 id 时由数据库回取）；
    若数据库未配置 id_field 且使用了 "id" 键，则退化为按总排号列查询（仅向后兼容，会告警）。
    """
    if not id_keys:
        return {}

    result: dict[str, dict[str, Any | None]] = {
        idt: {"param": None, "sn": None} for idt, _ in id_keys
    }

    src = _build_source_table(db_cfg)
    sn_col = db_cfg["id_column"]
    param_col = db_cfg["param_column"]
    id_field = db_cfg.get("id_field")
    sn_c = src.c[sn_col]
    param_c = src.c[param_col]

    sn_items = [idt for idt, k in id_keys if k == "sn"]
    id_items = [idt for idt, k in id_keys if k == "id"]

    engine = get_engine(db_cfg)

    def _fill_sn() -> None:
        if not sn_items:
            return
        with engine.connect() as conn:
            for chunk in _chunked(sn_items):
                stmt = select(sn_c, param_c).where(sn_c.in_(chunk))
                for row in conn.execute(stmt):
                    ident = str(row[0])
                    entry = result.get(ident)
                    if entry is not None:
                        entry["param"] = row[1]
                        entry["sn"] = ident

    def _fill_id() -> None:
        if not id_items:
            return
        if not id_field:
            # 未配置 id_field：降级为按总排号列查询
            logger.warning(
                "未配置 database.id_field，--id 将退化为按总排号列 [%s] 查询", sn_col
            )
            with engine.connect() as conn:
                for chunk in _chunked(id_items):
                    stmt = select(sn_c, param_c).where(sn_c.in_(chunk))
                    for row in conn.execute(stmt):
                        ident = str(row[0])
                        entry = result.get(ident)
                        if entry is not None:
                            entry["param"] = row[1]
                            entry["sn"] = ident
            return

        id_c = src.c[id_field]
        # 源表 id_field 为整型，将字符串标识符转为整数；无法转换的跳过并告警
        id_int_map: dict[int, str] = {}
        id_ints: list[int] = []
        for idt in id_items:
            try:
                v = int(idt)
            except ValueError:
                logger.warning("id 键 %s 无法转为整数（源表 %s 为整型），已跳过", idt, id_field)
                continue
            id_int_map[v] = idt
            id_ints.append(v)
        if not id_ints:
            return

        with engine.connect() as conn:
            for chunk in _chunked(id_ints):
                stmt = select(id_c, sn_c, param_c).where(id_c.in_(chunk))
                for row in conn.execute(stmt):
                    idv = row[0]
                    orig = id_int_map.get(idv)
                    if orig is None:
                        continue
                    entry = result.get(orig)
                    if entry is not None:
                        entry["param"] = row[2]
                        entry["sn"] = row[1]

    try:
        _fill_sn()
        _fill_id()
    except SQLAlchemyError as e:
        raise DatabaseError(f"数据库查询失败: {e}") from e

    return result


def fetch_param_by_id(
    db_cfg: dict[str, Any], identifier: str, key: str = "sn"
) -> dict[str, Any | None] | None:
    """单个查询的便捷封装，返回 {param, sn} 或 None（未查到）。"""
    return fetch_params_by_ids(db_cfg, [(identifier, key)]).get(identifier)


def fetch_all_ids(
    db_cfg: dict[str, Any],
    limit: int | None = None,
    order: str = "asc",
    id_min: int | None = None,
    id_max: int | None = None,
) -> list[str]:
    """拉取源表总排号列表，供全量/范围分类写入。

    排序按数据库真实 ID 列（config.id_field）；未配置 id_field 时降级为按总排号
    列排序并告警。order 取 "asc" / "desc"。id_min / id_max 按真实 ID 列做闭区间
    过滤（含端点）；任一为 None 表示不限制该侧。limit 取排序+过滤后的前 N 条。

    返回的仍是总排号列表，可直接喂给 classify_batch（键类型统一为 sn）。
    distinct + order_by + limit/offset 由 SQLAlchemy 按方言生成正确分页语法
    （PG: LIMIT n OFFSET 0；MSSQL: OFFSET 0 ROWS FETCH NEXT n ROWS ONLY）。
    """
    src = _build_source_table(db_cfg)
    sn_col = db_cfg["id_column"]
    id_field = db_cfg.get("id_field")
    sn_c = src.c[sn_col]

    engine = get_engine(db_cfg)

    order_dir = "desc" if order == "desc" else "asc"
    if id_field:
        id_c = src.c[id_field]
        stmt = select(sn_c.label("sn"), id_c.label("id"))
        wheres = [sn_c.isnot(None), id_c.isnot(None)]
        if id_min is not None:
            wheres.append(id_c >= int(id_min))
        if id_max is not None:
            wheres.append(id_c <= int(id_max))
        order_c = id_c
    else:
        logger.warning(
            "未配置 database.id_field，--limit/--range 将按总排号列排序（无法按真实 ID 排序/范围过滤）"
        )
        stmt = select(sn_c.label("sn"))
        wheres = [sn_c.isnot(None)]
        order_c = sn_c

    stmt = stmt.where(*wheres)
    stmt = stmt.order_by(order_c.asc() if order_dir == "asc" else order_c.desc())
    stmt = stmt.distinct()
    if limit is not None:
        stmt = stmt.limit(int(limit)).offset(0)

    try:
        with engine.connect() as conn:
            return [str(row[0]) for row in conn.execute(stmt)]
    except SQLAlchemyError as e:
        raise DatabaseError(f"查询总排号失败: {e}") from e


def fetch_existing_attachment_sns(db_cfg: dict[str, Any]) -> set[str]:
    """读取 Common.Attachment 中已存在的全部 SN（distinct）。

    用于"追加"模式做差集：源表总排号集合 减去 这里返回的 SN 集合，
    得到"源表有、但目标表还没写过"的总排号，只对这部分分类写入。

    返回 set[str]（空表返回空集合）；SN 统一转成 str，便于与源表总排号
    （同样是 str）直接做集合比较。NULL 值会被过滤掉。
    """
    engine = get_engine(db_cfg)
    try:
        with engine.connect() as conn:
            stmt = select(Attachment.SN).distinct()
            return {
                str(row[0]) for row in conn.execute(stmt) if row[0] is not None
            }
    except SQLAlchemyError as e:
        raise DatabaseError(f"读取已有附件 SN 失败: {e}") from e


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

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

    DELETE 的 IN 列表按 _PARAM_CHUNK 分片执行，避免 SN 过多时单条语句超过
    SQL Server 的 2100 绑定参数上限；整个"删 + 插"在同一事务内完成，语义不变。

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

    distinct_sn = sorted({r[0] for r in rows})
    engine = get_engine(db_cfg)

    try:
        with Session(engine) as session:
            # 分片 DELETE：单条 IN (...) 最多 _PARAM_CHUNK 个参数，避免触发 2100 上限
            deleted = 0
            for chunk in _chunked(distinct_sn):
                del_res = session.execute(
                    delete(Attachment).where(Attachment.SN.in_(chunk))
                )
                if del_res.rowcount:
                    deleted += del_res.rowcount
            session.execute(
                insert(Attachment),
                [
                    {"SN": r[0], "MajorCategory": r[1], "MinorCategory": r[2]}
                    for r in rows
                ],
            )
            inserted = len(rows)
            session.commit()
            logger.info(
                "upsert_attachments: 删除 %d 行, 插入 %d 行, 涉及 %d 个 SN",
                deleted, inserted, len(distinct_sn),
            )
            return (deleted, inserted)
    except SQLAlchemyError as e:
        raise DatabaseError(f"写入附件表失败: {e}") from e