a014a945af
- 新增 orm.py：按 db_type 构建引擎（postgresql+psycopg2 / mssql+pyodbc）， 声明式 Attachment 模型，init_schema 幂等建表 - 重写 db.py 为 SQLAlchemy Core 实现（动态 Table + quote 跨库正确引用、 id/sn 双键回查、跨库分页、先删后插幂等），对外签名不变 - 配置：config.yaml 默认 PostgreSQL，新增 config.mssql.yaml 保留 SQL Server， config_loader 支持可选 driver 与 db_type - write_attachments.py 新增 --init-db - 依赖 requirements.txt 增 sqlalchemy / psycopg2-binary（保留 pyodbc） - README 对齐：修正目标表字段描述并新增数据库抽象层小节
270 lines
9.4 KiB
Python
270 lines
9.4 KiB
Python
# -*- 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
|