#!/usr/bin/env python3 # -*- coding: utf-8 -*- """把通用图片表格 OCR 合并结果同步到 PostgreSQL 单表。""" from __future__ import annotations import argparse import csv import json import os import re import subprocess import tempfile import unicodedata from pathlib import Path from typing import Any DEFAULT_INPUT = Path("数据处理结果区/合并_图片表格_结构化.json") DEFAULT_SCHEMA = Path("数据处理工作区/06_PostgreSQL建表结构.sql") TEMPLATE_SCHEMA = Path("数据处理工作区/06_PostgreSQL建表结构.template.sql") CSV_FIELDS = [ "batch_name", "source_folder", "category_1", "category_2", "image_path", "image_name", "image_row_no", "unique_key", "record_data", "review_status", "review_notes", "manual_corrected", ] def normalize_text(value: Any) -> str: if value is None: return "" text = unicodedata.normalize("NFKC", str(value)).replace("\u3000", " ") return re.sub(r"\s+", " ", text).strip() def read_json(path: Path) -> Any: return json.loads(path.read_text(encoding="utf-8")) def scalar(value: Any) -> str: if value is None: return "" if isinstance(value, bool): return "true" if value else "false" return str(value) def sql_file(path: Path) -> Path: if path.exists(): return path if TEMPLATE_SCHEMA.exists(): return TEMPLATE_SCHEMA raise FileNotFoundError(f"找不到 SQL 文件: {path}") def quote_file(path: Path) -> str: return "'" + str(path.resolve()).replace("'", "''") + "'" def quote_table_name(name: str) -> str: name = normalize_text(name) if name.startswith('"') and name.endswith('"'): return name if re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", name): return f'"{name}"' raise ValueError(f"表名需要是简单标识符或已加双引号: {name}") def flatten_record(record: dict[str, Any], unique_key_field: str) -> dict[str, Any]: image = record.get("图片信息", {}) review = record.get("复核", {}) record_info = record.get("记录信息", {}) unique_key = normalize_text(record_info.get(unique_key_field, "")) return { "batch_name": record.get("处理批次", ""), "source_folder": record.get("来源文件夹", ""), "category_1": record.get("业务分类1", ""), "category_2": record.get("业务分类2", ""), "image_path": image.get("图片路径", ""), "image_name": image.get("图片名", ""), "image_row_no": image.get("图片内行号", ""), "unique_key": unique_key, "record_data": json.dumps(record_info, ensure_ascii=False, separators=(",", ":")), "review_status": review.get("状态", ""), "review_notes": "；".join(str(item) for item in review.get("提示", [])), "manual_corrected": bool(review.get("人工修正")), } def load_rows(input_path: Path, allow_empty_key: bool) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: archive = read_json(input_path) unique_key = normalize_text(archive.get("任务配置", {}).get("unique_key") or archive.get("汇总", {}).get("主唯一键")) if not unique_key: raise ValueError("合并结果中没有主唯一键配置，无法同步到默认 PostgreSQL 模板") rows: list[dict[str, Any]] = [] skipped: list[dict[str, Any]] = [] for record in archive.get("图片表格记录", []): row = flatten_record(record, unique_key) if not row["unique_key"] and not allow_empty_key: skipped.append(row) continue rows.append(row) return rows, skipped def write_temp_csv(rows: list[dict[str, Any]]) -> Path: temp_file = tempfile.NamedTemporaryFile( mode="w", encoding="utf-8", newline="", suffix="_image_table_records.csv", delete=False, ) with temp_file: writer = csv.DictWriter(temp_file, fieldnames=CSV_FIELDS) writer.writeheader() for row in rows: writer.writerow({field: scalar(row.get(field, "")) for field in CSV_FIELDS}) return Path(temp_file.name) def sync_to_postgres(args: argparse.Namespace, csv_path: Path, row_count: int) -> None: table_name = quote_table_name(args.table) truncate_sql = "" if args.append else f"TRUNCATE TABLE {table_name} RESTART IDENTITY;\n" sql = f"""\\set ON_ERROR_STOP on \\i {quote_file(sql_file(Path(args.schema)))} {truncate_sql}\\copy {table_name}({','.join(CSV_FIELDS)}) FROM {quote_file(csv_path)} WITH (FORMAT csv, HEADER true, NULL '') """ env = os.environ.copy() if args.password: env["PGPASSWORD"] = args.password elif env.get("WORKFLOW_DB_PASSWORD"): env["PGPASSWORD"] = env["WORKFLOW_DB_PASSWORD"] command = [ "psql", "-h", args.host, "-p", str(args.port), "-U", args.user, "-d", args.dbname, "-v", "ON_ERROR_STOP=1", ] completed = subprocess.run(command, input=sql, text=True, env=env, capture_output=True) print(completed.stdout, end="") if completed.returncode != 0: print(completed.stderr, end="") completed.check_returncode() print(json.dumps({"synced_table": args.table, "records": row_count}, ensure_ascii=False)) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--input", default=str(DEFAULT_INPUT), help="合并后的结构化 JSON") parser.add_argument("--schema", default=str(DEFAULT_SCHEMA), help="建表 SQL") parser.add_argument("--table", default=os.getenv("WORKFLOW_DB_TABLE", "Image_Table_Records"), help="目标表名") parser.add_argument("--host", default=os.getenv("WORKFLOW_DB_HOST", "127.0.0.1")) parser.add_argument("--port", default=os.getenv("WORKFLOW_DB_PORT", "5432")) parser.add_argument("--dbname", default=os.getenv("WORKFLOW_DB_NAME", "postgres")) parser.add_argument("--user", default=os.getenv("WORKFLOW_DB_USER", "postgres")) parser.add_argument("--password", default="", help="数据库密码；建议用 WORKFLOW_DB_PASSWORD") parser.add_argument("--append", action="store_true", help="追加导入，不清空表") parser.add_argument("--allow-empty-key", action="store_true", help="允许空 unique_key；默认跳过") return parser.parse_args() def main() -> None: args = parse_args() rows, skipped = load_rows(Path(args.input), args.allow_empty_key) csv_path = write_temp_csv(rows) try: sync_to_postgres(args, csv_path, len(rows)) if skipped: print(json.dumps({"skipped_empty_unique_key": len(skipped), "examples": skipped[:10]}, ensure_ascii=False)) finally: csv_path.unlink(missing_ok=True) if __name__ == "__main__": main()