# -*- 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