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:
2
.gitignore
vendored
2
.gitignore
vendored
@@ -23,6 +23,8 @@ logs/
|
||||
.env.local
|
||||
.env.*.local
|
||||
config.yaml
|
||||
config.*.yaml
|
||||
config.local.yaml
|
||||
|
||||
# Test & type caches
|
||||
.pytest_cache/
|
||||
|
||||
76
README.md
76
README.md
@@ -1,16 +1,23 @@
|
||||
# 布莱迪压力表 - 订单附件识别工具
|
||||
|
||||
根据"总排号"从 SQL Server 查询"新参数"字段,调用大语言模型判断该订单是否携带
|
||||
根据"总排号"从数据库查询"新参数"字段,调用大语言模型判断该订单是否携带
|
||||
附件,并以 JSON 输出结果。支持粗分类(资料/配件/耗材)和精分类("大类:细分
|
||||
类目",具体到针型阀、说明书等)两种粒度。
|
||||
|
||||
> 数据访问层基于 **SQLAlchemy** 抽象,支持 **SQL Server(mssql+pyodbc)** 与
|
||||
> **PostgreSQL(postgresql+psycopg2)** 两种数据库,通过 `config.database.db_type`
|
||||
> 切换。源表/目标表的标识符引用由 SQLAlchemy 按方言自动处理(PG 用 `"名"`,
|
||||
> MSSQL 用 `[名]`),无需改代码。
|
||||
|
||||
## 安装
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
`pyodbc` 需要系统已安装对应的 ODBC 驱动。本项目 `config.yaml` 中使用的版本为
|
||||
依赖包含 `sqlalchemy`、`psycopg2-binary`(PostgreSQL 驱动,已自带 libpq)与
|
||||
`pyodbc`(SQL Server 驱动)。使用 **PostgreSQL** 无需额外系统组件;使用
|
||||
**SQL Server** 需要系统已安装对应的 ODBC 驱动,本项目 `config.yaml` 中使用的版本为
|
||||
"ODBC Driver 18 for SQL Server",请按服务器实际安装的驱动版本填写 `database.driver`。
|
||||
若服务器上已有 SQL Server 管理工具/客户端环境,通常已包含该驱动;否则需自行
|
||||
安装 Microsoft 官方 ODBC Driver。
|
||||
@@ -20,9 +27,17 @@ pip install -r requirements.txt
|
||||
编辑 `config.yaml`,填入以下三部分(首次使用需替换为真实值,**请勿将含真实
|
||||
数据库密码 / API Key 的配置文件提交到版本库**):
|
||||
|
||||
- `database`:SQL Server 连接信息(含 `schema`)、表名、字段名。其中
|
||||
`id_column` 为**总排号**列(接口 `--sn` 使用),`id_field` 为数据库**真实 ID**
|
||||
列(接口 `--id` 使用);二者需按源表实际列名填写
|
||||
- `database`:数据库连接信息与源表/字段名。
|
||||
- `db_type`:数据库类型,`postgresql` 或 `mssql`(缺省按 `driver` 是否含
|
||||
"SQL Server" 推断;PostgreSQL 无需 `driver`)。
|
||||
- `server` / `port` / `database` / `username` / `password`:连接信息(两种库通用)。
|
||||
- `schema` / `table` / `id_column` / `id_field` / `param_column`:源表的
|
||||
schema、表名与列名(`id_column` 为**总排号**列,对应 `--sn`;`id_field` 为
|
||||
数据库**真实 ID** 列,对应 `--id`;`param_column` 为**新参数**列)。这些名称
|
||||
按源表实际填写,SQL Server / PostgreSQL 两端保持一致即可。
|
||||
- 仅 `mssql` 需要:`driver`、`trust_server_certificate`。
|
||||
- 切换到 SQL Server 时可直接用 `--config config.mssql.yaml`(已内置原 SQL Server
|
||||
连接信息)。
|
||||
- `llm`:OpenAI 兼容接口的 `base_url`、`api_key`、`model`
|
||||
- `business`:并发数、日志级别、默认分类模式(`default_mode`)、日志目录(`log_dir`)、
|
||||
是否启用"其他"兜底类目(`enable_other_category`,默认关闭)
|
||||
@@ -269,6 +284,10 @@ python write_attachments.py --mode coarse
|
||||
# 先预览将写入/跳过的行,不真正落库
|
||||
python write_attachments.py --dry-run
|
||||
|
||||
# 首次在目标库建表(幂等:仅创建 Common schema 与 Attachment 表,已存在则跳过;
|
||||
# 切换数据库或新环境首次部署前先执行一次)
|
||||
python write_attachments.py --init-db
|
||||
|
||||
# 允许"其他"兜底类目
|
||||
python write_attachments.py --mode fine --enable-other
|
||||
```
|
||||
@@ -285,22 +304,53 @@ python write_attachments.py --mode fine --enable-other
|
||||
`(SN, '无附件', '无')`,下游用 `WHERE MajorCategory <> '无附件'` 取真实附件。
|
||||
- **无法确定/失败**(`has_attachment=null`,含 `not_found` / `llm_*_error` /
|
||||
`db_error`):一律不写,既不当作无附件,也不留脏数据。
|
||||
- **幂等**:写入时对同一总排号先删除旧行再插入本次结果(依赖 `SN, MajorCategory,
|
||||
MinorCategory` 唯一索引),重跑安全。
|
||||
- **幂等**:写入时对同一总排号先删除旧行再插入本次结果(依赖 `(SN, MajorCategory,
|
||||
MinorCategory)` 复合主键保证唯一),重跑安全。
|
||||
|
||||
目标表 `Common.Attachment` 字段:`SN`(nvarchar(30),总排号/关联键)、
|
||||
`MajorCategory`(nvarchar(40),附件大类)、`MinorCategory`(nvarchar(40),附件小类),
|
||||
三者均 `NOT NULL`。
|
||||
目标表 `Common.Attachment` 字段(三列**复合主键**、均 `NOT NULL`;列类型由 ORM 模型
|
||||
`String(30)` / `String(40)` 按方言统一生成:PostgreSQL 端为 `varchar`,SQL Server 端为
|
||||
`nvarchar`):
|
||||
- `SN`:总排号/关联键
|
||||
- `MajorCategory`:附件大类
|
||||
- `MinorCategory`:附件小类
|
||||
|
||||
## 数据库抽象层(多库支持)
|
||||
|
||||
数据访问层已重构为 SQLAlchemy,由两层组成,业务代码(`classifier.py` /
|
||||
`write_attachments.py`)只调用 `db.py` 的 4 个函数,无需感知底层方言:
|
||||
|
||||
- `orm.py`:方言无关的底层。
|
||||
- `build_engine` / `get_engine`:按 `config.database.db_type` 生成 SQLAlchemy Engine
|
||||
(`postgresql+psycopg2` 或 `mssql+pyodbc`),含连接池复用(`pool_pre_ping`)与
|
||||
登录/语句超时;`get_engine` 按连接信息缓存 Engine,避免重复建池。
|
||||
- `Attachment`:目标表 `Common.Attachment` 的声明式 ORM 模型(三列复合主键,
|
||||
`quote=True` 保留大小写),跨库统一的建表/读写入口。
|
||||
- `init_schema`:方言感知地 `CREATE SCHEMA IF NOT EXISTS "Common"` + `create_all`,
|
||||
供 `--init-db` 幂等建表。
|
||||
- `db.py`:基于 SQLAlchemy Core 的查询/落库实现。
|
||||
- 源表(表名/列名含中文、由配置驱动)用动态 `Table(..., quote=True,
|
||||
quote_schema=True)` 构造,标识符引用由 SQLAlchemy 按方言生成(PG 用 `"名"`、
|
||||
MSSQL 用 `[名]`),彻底摆脱手写引号拼接。
|
||||
- 对外 4 个函数签名与旧版完全一致:`fetch_params_by_ids`(支持 `--id` 整型真实 ID
|
||||
与 `--sn` 总排号双键,并回查总排号)、`fetch_param_by_id`、`fetch_all_ids`
|
||||
(`distinct()` + `order_by()` + `limit()/offset(0)`,分页语法跨库自动适配)、
|
||||
`upsert_attachments`(先删后插,依赖复合主键幂等)。
|
||||
|
||||
**切换数据库**:默认 `config.yaml` 指向 PostgreSQL;切回 SQL Server 只需
|
||||
`python write_attachments.py --config config.mssql.yaml`(或 `main.py --config ...`)。
|
||||
两库源表 schema/表名/列名一致,仅需改连接信息与 `db_type`,无需改代码。
|
||||
|
||||
## 项目结构
|
||||
|
||||
```
|
||||
├── config.yaml # 配置文件
|
||||
├── config.yaml # 配置文件(默认 PostgreSQL;含 db_type 切换)
|
||||
├── config.mssql.yaml # SQL Server 版配置(数据库不可达时切换用,--config 指定)
|
||||
├── main.py # 命令行入口:分类并输出 JSON Lines(支持 --summary 在 stderr 打印批汇总;直接 import 同目录各模块)
|
||||
├── write_attachments.py # 写入入口:分类结果落库到 Common.Attachment(全表扫描支持 --limit/--order/--range 按真实 ID 排序与范围过滤;--id/--sn/--ids-file/--mode/--dry-run/--enable-other/--config;运行后打印 token/缓存汇总)
|
||||
├── write_attachments.py # 写入入口:分类结果落库到 Common.Attachment(全表扫描支持 --limit/--order/--range 按真实 ID 排序与范围过滤;--id/--sn/--ids-file/--mode/--dry-run/--enable-other/--init-db/--config;运行后打印 token/缓存汇总)
|
||||
├── requirements.txt
|
||||
├── config_loader.py # YAML 配置读取与校验
|
||||
├── db.py # SQL Server 查询与落库(fetch_params_by_ids 支持按 id/sn 双键查询、fetch_all_ids / upsert_attachments,pyodbc)
|
||||
├── orm.py # 数据库抽象层:SQLAlchemy 引擎构建(mssql/postgresql)、Attachment ORM 模型、init_schema 建表
|
||||
├── db.py # 基于 SQLAlchemy 的查询与落库(fetch_params_by_ids 支持按 id/sn 双键查询、fetch_all_ids / upsert_attachments,跨库无关)
|
||||
├── prompts.py # 提示词与细分类目枚举(唯一需要改分类边界时编辑的文件)
|
||||
├── llm_client.py # LLM 调用 (OpenAI 兼容接口)
|
||||
├── parser.py # LLM 输出格式校验与清洗 → 结构化数据
|
||||
|
||||
@@ -10,7 +10,7 @@ import yaml
|
||||
|
||||
_REQUIRED_KEYS = {
|
||||
"database": [
|
||||
"driver", "server", "port", "database", "schema", "username", "password",
|
||||
"server", "port", "database", "schema", "username", "password",
|
||||
"table", "id_column", "param_column", "connect_timeout", "query_timeout",
|
||||
],
|
||||
"llm": [
|
||||
|
||||
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):
|
||||
"""数据库连接或查询失败。"""
|
||||
|
||||
|
||||
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']};",
|
||||
# 兼容旧导入:把常量与异常从 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",
|
||||
]
|
||||
# 数据库服务器使用自签名/不受信任证书时,跳过证书链校验(连接仍保持加密)。
|
||||
# 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 _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(
|
||||
@@ -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:
|
||||
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})"
|
||||
)
|
||||
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()
|
||||
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)
|
||||
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"] = param
|
||||
entry["sn"] = sn_val
|
||||
except pyodbc.Error as e:
|
||||
raise DatabaseError(f"数据库查询失败: {e}") from e
|
||||
entry["param"] = row[1]
|
||||
entry["sn"] = ident
|
||||
|
||||
# 总排号键:直接查 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
|
||||
|
||||
177
orm.py
Normal file
177
orm.py
Normal file
@@ -0,0 +1,177 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""数据库抽象层:基于 SQLAlchemy 的多数据库支持(MSSQL / PostgreSQL)。
|
||||
|
||||
职责:
|
||||
- 按 db_type 构建对应方言的 Engine(连接信息完全来自 config.database)。
|
||||
- 定义目标附件表 Common.Attachment 的 ORM 模型(跨库统一的建表/读写入口)。
|
||||
- 提供幂等建库建表 init_schema(供 --init-db 调用)。
|
||||
|
||||
源表(productionContractData 等)列名由配置驱动、含中文/动态表名,不在这里建
|
||||
声明式模型,而是在 db.py 里用 Core 的 table()/column() + quote=True 动态构造,
|
||||
由 SQLAlchemy 负责生成各方言下正确的标识符引用(PG 用 "名",MSSQL 用 [名])。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import urllib.parse
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import (
|
||||
Column,
|
||||
Integer,
|
||||
String,
|
||||
create_engine,
|
||||
event,
|
||||
text,
|
||||
)
|
||||
from sqlalchemy.engine import Engine
|
||||
from sqlalchemy.orm import declarative_base
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DatabaseError(Exception):
|
||||
"""数据库连接或查询失败(统一异常,供上层 classify_batch / CLI 捕获)。"""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 目标附件表 Common.Attachment 的 ORM 模型
|
||||
# ---------------------------------------------------------------------------
|
||||
# 与项目记忆 MEMORY.md 一致:严格 3 字段,全部 NOT NULL;
|
||||
# 唯一约束 (SN, MajorCategory, MinorCategory)。这里用三列复合主键表达同一语义
|
||||
# (复合主键天然唯一,且无需额外唯一索引名),并强制 quote=True 保留大小写,
|
||||
# 使 PG 下表/列名保持 "Common"."Attachment" / "SN" 等原样,与 SQL Server 端一致。
|
||||
ATTACHMENT_SCHEMA = "Common"
|
||||
ATTACHMENT_TABLE = "Attachment"
|
||||
|
||||
SENTINEL_MAJOR = "无附件" # 无附件哨兵:大类固定写 '无附件'
|
||||
SENTINEL_MINOR = "无" # 哨兵/无小类占位:小类统一填 '无'
|
||||
NO_MINOR = "无" # coarse 真实附件 / 哨兵行 的小类占位
|
||||
|
||||
Base = declarative_base()
|
||||
|
||||
|
||||
class Attachment(Base):
|
||||
__tablename__ = ATTACHMENT_TABLE
|
||||
__table_args__ = (
|
||||
{"schema": ATTACHMENT_SCHEMA, "quote": True, "quote_schema": True},
|
||||
)
|
||||
|
||||
SN = Column("SN", String(30), primary_key=True, nullable=False, quote=True)
|
||||
MajorCategory = Column("MajorCategory", String(40), primary_key=True, nullable=False, quote=True)
|
||||
MinorCategory = Column("MinorCategory", String(40), primary_key=True, nullable=False, quote=True)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 连接信息解析
|
||||
# ---------------------------------------------------------------------------
|
||||
def resolve_db_type(db_cfg: dict[str, Any]) -> str:
|
||||
"""推断数据库类型:优先用显式 db_type,否则按 driver 是否含 'SQL Server' 判断。"""
|
||||
db_type = (db_cfg.get("db_type") or "").strip().lower()
|
||||
if db_type in ("mssql", "postgresql", "pgsql", "postgres"):
|
||||
return "postgresql" if db_type in ("postgresql", "pgsql", "postgres") else "mssql"
|
||||
driver = (db_cfg.get("driver") or "").lower()
|
||||
if "sql server" in driver:
|
||||
return "mssql"
|
||||
# 默认当作 PostgreSQL(本环境新引入的目标库)
|
||||
return "postgresql"
|
||||
|
||||
|
||||
def _build_url(db_cfg: dict[str, Any]) -> str:
|
||||
db_type = resolve_db_type(db_cfg)
|
||||
user = urllib.parse.quote_plus(db_cfg["username"])
|
||||
pw = urllib.parse.quote_plus(db_cfg["password"])
|
||||
host = db_cfg["server"]
|
||||
port = db_cfg["port"]
|
||||
database = db_cfg["database"]
|
||||
|
||||
if db_type == "postgresql":
|
||||
return f"postgresql+psycopg2://{user}:{pw}@{host}:{port}/{database}"
|
||||
# mssql
|
||||
driver = db_cfg.get("driver", "ODBC Driver 18 for SQL Server")
|
||||
drv = urllib.parse.quote_plus(driver)
|
||||
tsc = "yes" if db_cfg.get("trust_server_certificate", False) else "no"
|
||||
return (
|
||||
f"mssql+pyodbc://{user}:{pw}@{host}:{port}/{database}"
|
||||
f"?driver={drv}&TrustServerCertificate={tsc}"
|
||||
)
|
||||
|
||||
|
||||
def build_engine(db_cfg: dict[str, Any], echo: bool = False) -> Engine:
|
||||
"""按配置构建 SQLAlchemy Engine。
|
||||
|
||||
- 登录超时:mssql 通过 connect_args.timeout(pyodbc 登录超时);pg 由驱动处理。
|
||||
- 语句超时:通过 connect 事件设置(mssql: dbapi_conn.timeout;pg: SET statement_timeout)。
|
||||
- pool_pre_ping 开启,避免跨长时间空闲的连接失效。
|
||||
"""
|
||||
url = _build_url(db_cfg)
|
||||
db_type = resolve_db_type(db_cfg)
|
||||
connect_timeout = int(db_cfg.get("connect_timeout", 10))
|
||||
query_timeout = int(db_cfg.get("query_timeout", 15))
|
||||
|
||||
connect_args: dict[str, Any] = {}
|
||||
if db_type == "mssql":
|
||||
connect_args["timeout"] = connect_timeout # pyodbc 登录超时
|
||||
|
||||
try:
|
||||
engine = create_engine(url, connect_args=connect_args, pool_pre_ping=True, future=True, echo=echo)
|
||||
except Exception as e: # pragma: no cover - 配置/驱动错误
|
||||
raise DatabaseError(f"创建数据库引擎失败: {e}") from e
|
||||
|
||||
if db_type == "mssql":
|
||||
@event.listens_for(engine, "connect")
|
||||
def _set_mssql_timeout(dbapi_conn, _rec) -> None:
|
||||
try:
|
||||
dbapi_conn.timeout = query_timeout # pyodbc 语句超时(秒)
|
||||
except Exception: # pragma: no cover
|
||||
pass
|
||||
else:
|
||||
@event.listens_for(engine, "connect")
|
||||
def _set_pg_timeout(dbapi_conn, _rec) -> None:
|
||||
try:
|
||||
cur = dbapi_conn.cursor()
|
||||
cur.execute(f"SET statement_timeout = {query_timeout * 1000}")
|
||||
cur.close()
|
||||
except Exception: # pragma: no cover
|
||||
pass
|
||||
|
||||
return engine
|
||||
|
||||
|
||||
_ENGINES: dict[tuple, Engine] = {}
|
||||
|
||||
|
||||
def get_engine(db_cfg: dict[str, Any], echo: bool = False) -> Engine:
|
||||
"""获取(并缓存)Engine,避免每次调用都重建连接池。"""
|
||||
key = (
|
||||
resolve_db_type(db_cfg),
|
||||
db_cfg.get("server"),
|
||||
db_cfg.get("port"),
|
||||
db_cfg.get("database"),
|
||||
db_cfg.get("username"),
|
||||
)
|
||||
engine = _ENGINES.get(key)
|
||||
if engine is None:
|
||||
engine = build_engine(db_cfg, echo=echo)
|
||||
_ENGINES[key] = engine
|
||||
return engine
|
||||
|
||||
|
||||
def init_schema(engine: Engine) -> None:
|
||||
"""幂等初始化目标库:确保 Common schema 存在并建 Attachment 表。
|
||||
|
||||
源表(如 productionContractData)是既有数据,不在此处创建/修改。
|
||||
仅在目标库操作,用于首次迁移/部署时准备落库表。
|
||||
"""
|
||||
dialect = engine.dialect.name # 'postgresql' / 'mssql'
|
||||
with engine.begin() as conn:
|
||||
if dialect == "postgresql":
|
||||
conn.execute(text('CREATE SCHEMA IF NOT EXISTS "Common"'))
|
||||
else: # mssql
|
||||
conn.execute(text(
|
||||
"IF NOT EXISTS (SELECT 1 FROM sys.schemas WHERE name='Common') "
|
||||
"EXEC('CREATE SCHEMA [Common]')"
|
||||
))
|
||||
# 建表(checkfirst=True:已存在则跳过)
|
||||
Base.metadata.create_all(conn, checkfirst=True)
|
||||
logger.info("init_schema 完成:%s.%s 已就绪", ATTACHMENT_SCHEMA, ATTACHMENT_TABLE)
|
||||
@@ -1,3 +1,5 @@
|
||||
pyyaml>=6.0
|
||||
pyodbc>=5.0
|
||||
sqlalchemy>=2.0
|
||||
psycopg2-binary>=2.9
|
||||
openai>=1.0
|
||||
|
||||
@@ -49,6 +49,8 @@ from db import ( # noqa: E402
|
||||
SENTINEL_MINOR,
|
||||
NO_MINOR,
|
||||
fetch_all_ids,
|
||||
get_engine,
|
||||
init_schema,
|
||||
upsert_attachments,
|
||||
)
|
||||
|
||||
@@ -154,6 +156,10 @@ def main() -> None:
|
||||
"--dry-run", action="store_true",
|
||||
help="只打印将要写入/跳过的统计,不连接目标表写入",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--init-db", action="store_true",
|
||||
help="仅初始化目标库:幂等创建 Common.Attachment 表(含 Common schema),不进行分类/写入",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
try:
|
||||
@@ -169,6 +175,16 @@ def main() -> None:
|
||||
)
|
||||
logger = logging.getLogger("write_attachments")
|
||||
|
||||
# 仅初始化目标表(幂等建 Common schema + Attachment 表),然后退出
|
||||
if args.init_db:
|
||||
try:
|
||||
init_schema(get_engine(cfg["database"]))
|
||||
except DatabaseError as e:
|
||||
print(f"[init-db 失败] {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
print(f"[init-db] 已完成:目标表 {ATTACHMENT_SCHEMA}.{ATTACHMENT_TABLE} 已就绪(幂等)")
|
||||
return
|
||||
|
||||
mode = args.mode or cfg["business"]["default_mode"]
|
||||
log_dir = args.log_dir or cfg["business"]["log_dir"]
|
||||
# enable_other_category 是可选配置项,缺省按"关闭";--enable-other 只能从关到开
|
||||
|
||||
Reference in New Issue
Block a user