193 lines
6.7 KiB
Python
Executable File
193 lines
6.7 KiB
Python
Executable File
#!/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()
|