Files
attachment_classifier/config_loader.py
pengq a014a945af 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 对齐:修正目标表字段描述并新增数据库抽象层小节
2026-07-26 11:59:55 +08:00

89 lines
3.3 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# -*- coding: utf-8 -*-
"""配置文件加载与校验。"""
from __future__ import annotations
import sys
from pathlib import Path
from typing import Any
import yaml
_REQUIRED_KEYS = {
"database": [
"server", "port", "database", "schema", "username", "password",
"table", "id_column", "param_column", "connect_timeout", "query_timeout",
],
"llm": [
"base_url", "api_key", "model", "temperature", "max_tokens",
"disable_thinking", "timeout", "max_retry", "retry_backoff_seconds",
],
"business": ["max_workers", "log_level", "default_mode", "log_dir", "format_retry"],
}
_VALID_MODES = {"coarse", "fine"}
class ConfigError(Exception):
"""配置文件缺失或格式错误。"""
def load_config(path: str | Path = "config.yaml") -> dict[str, Any]:
"""读取 YAML 配置文件并做基础字段校验。
校验失败会抛出 ConfigError 并说明缺哪个字段,避免程序运行到一半才因为
配置缺失而报出难以理解的异常。
"""
p = Path(path)
if not p.exists():
raise ConfigError(f"配置文件不存在: {p.resolve()}")
with p.open("r", encoding="utf-8") as f:
try:
cfg = yaml.safe_load(f)
except yaml.YAMLError as e:
raise ConfigError(f"配置文件 YAML 解析失败: {e}") from e
if not isinstance(cfg, dict):
raise ConfigError("配置文件顶层必须是一个字典 (database/llm/business)")
missing_sections = [s for s in _REQUIRED_KEYS if s not in cfg]
if missing_sections:
raise ConfigError(f"配置文件缺少顶层配置块: {missing_sections}")
for section, keys in _REQUIRED_KEYS.items():
section_cfg = cfg[section]
if not isinstance(section_cfg, dict):
raise ConfigError(f"配置块 '{section}' 必须是字典")
missing = [k for k in keys if k not in section_cfg]
if missing:
raise ConfigError(f"配置块 '{section}' 缺少字段: {missing}")
default_mode = cfg["business"]["default_mode"]
if default_mode not in _VALID_MODES:
raise ConfigError(
f"business.default_mode 必须是 {_VALID_MODES} 之一,实际为: {default_mode!r}"
)
# enable_other_category 是后加的可选字段,不放进 _REQUIRED_KEYS——避免
# 旧配置文件因为没有这个字段就直接报"缺字段"。不写就按 False关闭处理
# 一旦写了就必须是布尔值,防止 "false"(字符串) 这类误写被当成真值静默生效。
if "enable_other_category" in cfg["business"]:
other_flag = cfg["business"]["enable_other_category"]
if not isinstance(other_flag, bool):
raise ConfigError(
f"business.enable_other_category 必须是布尔值 true/false实际为: {other_flag!r}"
)
# 简单的占位符提醒,防止用户忘记改配置直接运行
placeholders = []
if cfg["database"]["password"] == "CHANGE_ME":
placeholders.append("database.password")
if cfg["llm"]["api_key"] == "CHANGE_ME":
placeholders.append("llm.api_key")
if placeholders:
print(
f"[警告] 以下配置项仍是占位符,请修改 config.yaml: {placeholders}",
file=sys.stderr,
)
return cfg