70b6fb3767
此前 --limit 仅按总排号排序取前 N 个。现扩展为：
- 排序按数据库真实 ID 列（config.id_field），新增 --order {asc,desc}（默认 asc）
- 新增 --range START,END 按真实 ID 闭区间过滤（与 --order/--limit 可组合）
- 上述三参数仅全表扫描生效，与 --id/--sn/--ids-file 互斥
db.py: fetch_all_ids 新增 order/id_min/id_max 参数，用派生表
（内层 DISTINCT 取 (总排号,ID) 配对，外层按 ID 排序）规避
SQL Server "SELECT DISTINCT 时 ORDER BY 列须在选择列表" 的限制；
id_field 未配置时降级按总排号排序并告警。
write_attachments.py: 新增 --order/--range CLI 参数并接线，含 --range
格式校验（须为 START,END 两个整数）；模块 docstring 补示例。
README.md: 写入用法块与项目结构注释补充 --order/--range 说明。
自测（小数据量，只读+部分 dry-run）：--limit 3 升/降序、--range 800,805
升序、--range 800,805 --order desc --limit 3 均验证返回总排号按真实 ID
正确排序且在范围内；dry-run 不落库；--range 格式错误正确报错退出。
Co-Authored-By: WorkBuddy <workbuddy@tencent.com>
252 lines
10 KiB
Python
252 lines
10 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）：写哨兵行 (SN, '无附件', '无')
|
||
# - 无小类的行（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,
|
||
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 取 (总排号, ID) 配对，外层按 ID 排序）规避 SQL Server
|
||
"SELECT DISTINCT 时 ORDER BY 列须出现在选择列表"的限制。
|
||
"""
|
||
schema = db_cfg["schema"]
|
||
table = db_cfg["table"]
|
||
sn_col = db_cfg["id_column"] # 总排号列（返回列）
|
||
id_field = db_cfg.get("id_field") # 真实 ID 列（排序/过滤用）
|
||
qualified = f"[{schema}].[{table}]"
|
||
|
||
order_dir = "DESC" if order == "desc" else "ASC"
|
||
inner_where = [f"[{sn_col}] IS NOT NULL"]
|
||
inner_params: list[Any] = []
|
||
|
||
if id_field:
|
||
inner_where.append(f"[{id_field}] IS NOT NULL")
|
||
if id_min is not None:
|
||
inner_where.append(f"[{id_field}] >= ?")
|
||
inner_params.append(int(id_min))
|
||
if id_max is not None:
|
||
inner_where.append(f"[{id_field}] <= ?")
|
||
inner_params.append(int(id_max))
|
||
order_expr = "[id]" # 派生表别名
|
||
inner_select = f"SELECT DISTINCT [{sn_col}] AS sn, [{id_field}] AS id"
|
||
else:
|
||
logger.warning(
|
||
"未配置 database.id_field，--limit/--range 将按总排号列排序（无法按真实 ID 排序/范围过滤）"
|
||
)
|
||
order_expr = "[sn]"
|
||
inner_select = f"SELECT DISTINCT [{sn_col}] AS sn"
|
||
|
||
inner_sql = f"{inner_select} FROM {qualified} WHERE {' AND '.join(inner_where)}"
|
||
sql = f"SELECT [sn] FROM ({inner_sql}) AS t ORDER BY {order_expr} {order_dir}"
|
||
params = list(inner_params)
|
||
if limit is not None:
|
||
sql += " 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: 本次要写入的 (SN, MajorCategory, MinorCategory) 列表。
|
||
对出现的每个 SN 先 DELETE 其旧行，再批量 INSERT——保证重跑总是反映
|
||
最新分类结果，不会因唯一索引 (SN, 大类, 小类) 冲突而失败。
|
||
|
||
返回 (deleted_rows, inserted_rows) 计数。
|
||
"""
|
||
if not rows:
|
||
return (0, 0)
|
||
|
||
qualified = f"[{ATTACHMENT_SCHEMA}].[{ATTACHMENT_TABLE}]"
|
||
distinct_sn = 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) 删除本批所有 SN 的旧行
|
||
del_ph = ",".join("?" for _ in distinct_sn)
|
||
cursor.execute(
|
||
f"DELETE FROM {qualified} WHERE [SN] IN ({del_ph})", distinct_sn
|
||
)
|
||
deleted = cursor.rowcount
|
||
# 2) 插入新行
|
||
cursor.executemany(
|
||
f"INSERT INTO {qualified} ([SN], [MajorCategory], [MinorCategory]) "
|
||
f"VALUES (?, ?, ?)",
|
||
rows,
|
||
)
|
||
inserted = len(rows)
|
||
conn.commit()
|
||
logger.info(
|
||
"upsert_attachments: 删除 %d 行, 插入 %d 行, 涉及 %d 个 SN",
|
||
deleted, inserted, len(distinct_sn),
|
||
)
|
||
return (deleted, inserted)
|
||
except pyodbc.Error as e:
|
||
raise DatabaseError(f"写入附件表失败: {e}") from e
|