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:
300
db.py
300
db.py
@@ -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 0;MSSQL: 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
|
||||
|
||||
Reference in New Issue
Block a user