refactor: 引入 SQLAlchemy 多数据库抽象(PostgreSQL + SQL Server)

- 新增 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 对齐:修正目标表字段描述并新增数据库抽象层小节
This commit is contained in:
pengq
2026-07-26 11:59:55 +08:00
parent 70b6fb3767
commit a014a945af
7 changed files with 420 additions and 155 deletions

300
db.py
View File

@@ -1,34 +1,69 @@
# -*- coding: utf-8 -*-
"""SQL Server 数据访问层,通过 pyodbc 按总排号批量查询新参数字段。"""
"""数据访问层:基于 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
import pyodbc
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__)
class DatabaseError(Exception):
"""数据库连接或查询失败。"""
# 兼容旧导入:把常量与异常从 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_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']};",
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),
]
# 数据库服务器使用自签名/不受信任证书时,跳过证书链校验(连接仍保持加密)。
# ODBC Driver 17/18 默认 Encrypt=Yes遇到自签证书会报"不受信任的颁发机构"
# 加 TrustServerCertificate=Yes 即可信任该证书。设为 false 时不影响原有行为。
if db_cfg.get("trust_server_certificate", False):
parts.append("TrustServerCertificate=Yes;")
return "".join(parts)
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(
@@ -41,12 +76,9 @@ def fetch_params_by_ids(
"id" -> 按数据库真实 ID 列(配置 id_field查询并回取对应的总排号。
返回 dict{identifier: {"param": 新参数文本|None, "sn": 总排号|None}}。
标识符作为 key 原样保留(便于回查);未查到的标识符其 param/sn 为 None
"查到但内容为空"区分开。sn 为对应的总排号(键类型为 sn 时即标识符本身,
键类型为 id 时由数据库回取);若数据库未配置 id_field 且使用了 "id" 键,
则退化为按总排号列查询(仅向后兼容,会在日志告警)。
一对一关系:同一标识符出现多条记录时取第一条并记录 WARNING。
标识符作为 key 原样保留(便于回查);未查到的标识符其 param/sn 为 None
sn 为对应的总排号(键类型为 sn 时即标识符本身,键类型为 id 时由数据库回取);
若数据库未配置 id_field 且使用了 "id" 键,则退化为按总排号列查询(仅向后兼容,会告警)。
"""
if not id_keys:
return {}
@@ -55,69 +87,80 @@ def fetch_params_by_ids(
idt: {"param": None, "sn": None} for idt, _ in id_keys
}
schema = db_cfg["schema"]
table = db_cfg["table"]
src = _build_source_table(db_cfg)
sn_col = db_cfg["id_column"]
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}]"
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"]
def _fill(identifiers: list[str], where_col: str, is_id_key: bool) -> None:
if not identifiers:
engine = get_engine(db_cfg)
def _fill_sn() -> None:
if not sn_items:
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})"
)
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
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:
def _fill_id() -> None:
if not id_items:
return
if not id_field:
# 未配置 id_field降级为按总排号列查询
logger.warning(
"未配置 database.id_field--id 将退化为按总排号列 [%s] 查询", sn_col
)
_fill(id_items, sn_col, is_id_key=False)
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
@@ -129,20 +172,6 @@ def fetch_param_by_id(
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,
@@ -157,51 +186,44 @@ def fetch_all_ids(
过滤(含端点);任一为 None 表示不限制该侧。limit 取排序+过滤后的前 N 条。
返回的仍是总排号列表,可直接喂给 classify_batch键类型统一为 sn
用派生表(内层 DISTINCT 取 (总排号, ID) 配对,外层按 ID 排序)规避 SQL Server
"SELECT DISTINCT 时 ORDER BY 列须出现在选择列表"的限制
distinct + order_by + limit/offset 由 SQLAlchemy 按方言生成正确分页语法
PG: LIMIT n OFFSET 0MSSQL: OFFSET 0 ROWS FETCH NEXT n ROWS ONLY
"""
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}]"
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]
order_dir = "DESC" if order == "desc" else "ASC"
inner_where = [f"[{sn_col}] IS NOT NULL"]
inner_params: list[Any] = []
engine = get_engine(db_cfg)
order_dir = "desc" if order == "desc" else "asc"
if id_field:
inner_where.append(f"[{id_field}] IS NOT NULL")
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:
inner_where.append(f"[{id_field}] >= ?")
inner_params.append(int(id_min))
wheres.append(id_c >= 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"
wheres.append(id_c <= int(id_max))
order_c = id_c
else:
logger.warning(
"未配置 database.id_field--limit/--range 将按总排号列排序(无法按真实 ID 排序/范围过滤)"
)
order_expr = "[sn]"
inner_select = f"SELECT DISTINCT [{sn_col}] AS sn"
stmt = select(sn_c.label("sn"))
wheres = [sn_c.isnot(None)]
order_c = sn_c
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)
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:
sql += " OFFSET 0 ROWS FETCH NEXT ? ROWS ONLY"
params.append(int(limit))
stmt = stmt.limit(int(limit)).offset(0)
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:
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
@@ -213,39 +235,35 @@ def upsert_attachments(
rows: 本次要写入的 (SN, MajorCategory, MinorCategory) 列表。
对出现的每个 SN 先 DELETE 其旧行,再批量 INSERT——保证重跑总是反映
最新分类结果,不会因唯一索引 (SN, 大类, 小类) 冲突而失败。
最新分类结果,不会因唯一约束(复合主键 (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})
engine = get_engine(db_cfg)
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
with Session(engine) as session:
del_res = session.execute(
delete(Attachment).where(Attachment.SN.in_(distinct_sn))
)
deleted = cursor.rowcount
# 2) 插入新行
cursor.executemany(
f"INSERT INTO {qualified} ([SN], [MajorCategory], [MinorCategory]) "
f"VALUES (?, ?, ?)",
rows,
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)
conn.commit()
session.commit()
logger.info(
"upsert_attachments: 删除 %d 行, 插入 %d 行, 涉及 %d 个 SN",
deleted, inserted, len(distinct_sn),
)
return (deleted, inserted)
except pyodbc.Error as e:
except SQLAlchemyError as e:
raise DatabaseError(f"写入附件表失败: {e}") from e