# -*- 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", "upsert_attachments", "init_schema", ] 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 stmt = select(sn_c, param_c).where(sn_c.in_(sn_items)) with engine.connect() as conn: 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 ) stmt = select(sn_c, param_c).where(sn_c.in_(id_items)) with engine.connect() as conn: 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 stmt = select(id_c, sn_c, param_c).where(id_c.in_(id_ints)) with engine.connect() as conn: 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 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,大类,小类))冲突而失败。 返回 (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: del_res = session.execute( delete(Attachment).where(Attachment.SN.in_(distinct_sn)) ) deleted = del_res.rowcount if del_res.rowcount is not None else 0 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