diff --git a/.gitignore b/.gitignore index 63b9a20..3f9fa54 100644 --- a/.gitignore +++ b/.gitignore @@ -23,6 +23,8 @@ logs/ .env.local .env.*.local config.yaml +config.*.yaml +config.local.yaml # Test & type caches .pytest_cache/ diff --git a/README.md b/README.md index fbc289f..69139fc 100644 --- a/README.md +++ b/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 输出格式校验与清洗 → 结构化数据 diff --git a/config_loader.py b/config_loader.py index daf8500..aefb7c9 100644 --- a/config_loader.py +++ b/config_loader.py @@ -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": [ diff --git a/db.py b/db.py index 5ae72a6..65ff2f3 100644 --- a/db.py +++ b/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 diff --git a/orm.py b/orm.py new file mode 100644 index 0000000..ff34fe2 --- /dev/null +++ b/orm.py @@ -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) diff --git a/requirements.txt b/requirements.txt index b83aae8..01a47bf 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,5 @@ pyyaml>=6.0 pyodbc>=5.0 +sqlalchemy>=2.0 +psycopg2-binary>=2.9 openai>=1.0 diff --git a/write_attachments.py b/write_attachments.py index 0e26940..85a57cc 100644 --- a/write_attachments.py +++ b/write_attachments.py @@ -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 只能从关到开