ee566e3fbd
此前接口 --id 实际按总排号列查询，与数据库真实 ID 字段语义混淆。
现明确区分两种键类型：
- --id -> 数据库真实 ID 列（config id_field）
- --sn -> 总排号列（config id_column，即原先 --id 的语义）
- --ids-file 视为总排号（键类型 sn）
db.py: fetch_params_by_ids 改为接受 (标识符, 键类型) 列表，返回
{标识符: {param, sn}}；id 键回取对应总排号；新增 id_field 配置读取
classifier.py: classify_batch/classify_single 透传键类型，输出 zong_pai_hao
一律为回查到的总排号，not_found 时回退输入标识符便于追溯
main.py / write_attachments.py: 新增 --sn，--id 改指真实 ID，二者可混用；
全表扫描回退仍按总排号(sn)查询
README.md: 对齐 --id/--sn 语义并补充 id_field 配置说明
注：config.yaml 含密钥被 .gitignore 忽略，id_field 仅存于本地配置；
db.py 在未配置 id_field 时优雅降级为按总排号查询并告警。
Co-Authored-By: WorkBuddy <workbuddy@tencent.com>
216 lines
8.6 KiB
Python
216 lines
8.6 KiB
Python
# -*- 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], 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" 键，
|
||
则退化为按总排号列查询（仅向后兼容，会在日志告警）。
|
||
|
||
一对一关系：同一标识符出现多条记录时取第一条并记录 WARNING。
|
||
"""
|
||
if not id_keys:
|
||
return {}
|
||
|
||
result: dict[str, dict[str, Any | None]] = {
|
||
idt: {"param": None, "sn": None} for idt, _ in id_keys
|
||
}
|
||
|
||
schema = db_cfg["schema"]
|
||
table = db_cfg["table"]
|
||
param_col = db_cfg["param_column"]
|
||
sn_col = db_cfg["id_column"] # 总排号列
|
||
id_field = db_cfg.get("id_field") # 真实 ID 列，可缺省
|
||
qualified_table = f"[{schema}].[{table}]"
|
||
|
||
sn_items = [idt for idt, k in id_keys if k == "sn"]
|
||
id_items = [idt for idt, k in id_keys if k == "id"]
|
||
|
||
def _fill(identifiers: list[str], where_col: str, is_id_key: bool) -> None:
|
||
if not identifiers:
|
||
return
|
||
# 表名/schema 加中括号转义，不能写成 [schema.table]
|
||
placeholders = ",".join("?" for _ in identifiers)
|
||
if is_id_key:
|
||
# 查真实 ID 列，同时回取总排号列(sn_col)作为 sn
|
||
sql = (
|
||
f"SELECT [{where_col}], [{sn_col}], [{param_col}] "
|
||
f"FROM {qualified_table} WHERE [{where_col}] IN ({placeholders})"
|
||
)
|
||
else:
|
||
sql = (
|
||
f"SELECT [{where_col}], [{param_col}] "
|
||
f"FROM {qualified_table} WHERE [{where_col}] IN ({placeholders})"
|
||
)
|
||
|
||
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, identifiers)
|
||
seen: set[str] = set()
|
||
for row in cursor.fetchall():
|
||
if is_id_key:
|
||
ident, sn_val, param = row[0], row[1], row[2]
|
||
else:
|
||
ident, param = row[0], row[1]
|
||
sn_val = ident # 总排号即标识符本身
|
||
ident = str(ident)
|
||
if ident in seen:
|
||
logger.warning("标识符 %s 在列 [%s] 上重复，已取第一条", ident, where_col)
|
||
continue
|
||
seen.add(ident)
|
||
entry = result.get(ident)
|
||
if entry is not None:
|
||
entry["param"] = param
|
||
entry["sn"] = sn_val
|
||
except pyodbc.Error as e:
|
||
raise DatabaseError(f"数据库查询失败: {e}") from e
|
||
|
||
# 总排号键：直接查 id_column
|
||
_fill(sn_items, sn_col, is_id_key=False)
|
||
# 真实 ID 键：查 id_field；未配置时降级为总排号列并告警
|
||
if id_items:
|
||
if id_field:
|
||
_fill(id_items, id_field, is_id_key=True)
|
||
else:
|
||
logger.warning(
|
||
"未配置 database.id_field，--id 将退化为按总排号列 [%s] 查询", sn_col
|
||
)
|
||
_fill(id_items, sn_col, is_id_key=False)
|
||
|
||
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)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 附件分类落库（目标表 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
|