"""Import the supplied anomaly and judgment-standard workbooks. The import is intentionally idempotent: every source row gets a stable SHA-256 fingerprint based on its file, worksheet, and Excel row number. """ from __future__ import annotations import argparse import asyncio import hashlib import logging import re from collections.abc import Iterable, Sequence from datetime import date, datetime from pathlib import Path from typing import Any from app.core.config import get_settings from app.models.entities import AnomalyCase, StandardRule from openpyxl import load_workbook from sqlalchemy.dialects.postgresql import insert from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine LOGGER = logging.getLogger("dieya_risk_prediction.import") ANOMALY_SHEET = "异常总览（新流程分类）" STANDARD_SHEETS = ("包装", "包材", "法规", "内料") def clean_text(value: Any) -> str | None: """Normalize spreadsheet text while preserving meaningful content.""" if value is None: return None text = str(value).replace("\u200b", "").strip() return text or None def parse_date(value: Any) -> date | None: """Parse the date formats used by the source workbook.""" if isinstance(value, datetime): return value.date() if isinstance(value, date): return value text = clean_text(value) if not text: return None text = re.sub(r"\s+00:00:00$", "", text) text = text.replace("年", "-").replace("月", "-").replace("日", "") for fmt in ("%Y-%m-%d", "%Y/%m/%d", "%Y.%m.%d"): try: return datetime.strptime(text, fmt).date() except ValueError: continue month_only = re.fullmatch(r"(\d{4})[-/.](\d{1,2})", text) if month_only: return date(int(month_only.group(1)), int(month_only.group(2)), 1) return None def fingerprint(file_name: str, sheet_name: str, row_number: int) -> str: source = f"{file_name}\x1f{sheet_name}\x1f{row_number}" return hashlib.sha256(source.encode("utf-8")).hexdigest() def row_values(row: Sequence[Any], size: int) -> list[Any]: """Return a fixed-size row so trailing empty cells are handled uniformly.""" values = list(row) return values[:size] + [None] * max(0, size - len(values)) def read_anomaly_rows(path: Path) -> list[dict[str, Any]]: workbook = load_workbook(path, data_only=True) sheet = workbook[ANOMALY_SHEET] rows: list[dict[str, Any]] = [] for excel_row, raw_row in enumerate(sheet.iter_rows(min_row=2, values_only=True), start=2): values = row_values(raw_row, 22) rows.append( { "date": parse_date(values[1]), "product_code": clean_text(values[2]), "product_name": clean_text(values[3]), "brand": clean_text(values[4]), "specification": clean_text(values[5]), "stage_new": clean_text(values[6]), "stage_old": clean_text(values[7]), "anomaly_type": clean_text(values[8]), "material_type": clean_text(values[9]), "supplier": clean_text(values[10]), "factory": clean_text(values[11]), "person": clean_text(values[12]), "description": clean_text(values[13]), "root_cause": clean_text(values[14]), "solution": clean_text(values[15]), "result": clean_text(values[16]), "responsible": clean_text(values[17]), "defect_level": clean_text(values[18]) if clean_text(values[18]) != "/" else None, "is_resolved": clean_text(values[19]), "source_file": clean_text(values[20]) or path.name, "remark": clean_text(values[21]), "source_type": "Excel导入", "import_fingerprint": fingerprint(path.name, ANOMALY_SHEET, excel_row), } ) return rows def read_standard_rows(path: Path) -> list[dict[str, Any]]: workbook = load_workbook(path, data_only=True) rows: list[dict[str, Any]] = [] for sheet_name in STANDARD_SHEETS: sheet = workbook[sheet_name] for excel_row, raw_row in enumerate(sheet.iter_rows(min_row=2, values_only=True), start=2): values = row_values(raw_row, 9) if not any(value is not None and str(value).strip() for value in values): continue rows.append( { "category_lv1": clean_text(values[1]) or sheet_name, "category_lv2": clean_text(values[2]), "category_lv3": clean_text(values[3]), "risk_point": clean_text(values[4]) or "未填写", "judge_standard": clean_text(values[5]) or "未填写", "regulation": clean_text(values[6]), "attachment": " | ".join( value for value in (clean_text(values[7]), clean_text(values[8])) if value ) or None, "status": "启用", "import_fingerprint": fingerprint(path.name, sheet_name, excel_row), } ) return rows async def insert_ignore_duplicates( session: AsyncSession, model: type[AnomalyCase] | type[StandardRule], rows: Iterable[dict[str, Any]], ) -> tuple[int, int]: payload = list(rows) if not payload: return 0, 0 inserted = 0 batch_size = 100 for start in range(0, len(payload), batch_size): batch = payload[start : start + batch_size] statement = insert(model).values(batch).on_conflict_do_nothing(index_elements=["import_fingerprint"]) result = await session.execute(statement) inserted += result.rowcount or 0 return inserted, len(payload) - inserted async def import_workbooks(anomaly_path: Path, standard_path: Path) -> dict[str, int]: anomaly_rows = read_anomaly_rows(anomaly_path) standard_rows = read_standard_rows(standard_path) settings = get_settings() engine = create_async_engine(settings.database_url, pool_pre_ping=True) session_factory = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) try: async with session_factory() as session: anomaly_inserted, anomaly_skipped = await insert_ignore_duplicates(session, AnomalyCase, anomaly_rows) standard_inserted, standard_skipped = await insert_ignore_duplicates(session, StandardRule, standard_rows) await session.commit() finally: await engine.dispose() return { "anomaly_source_rows": len(anomaly_rows), "anomaly_inserted": anomaly_inserted, "anomaly_skipped": anomaly_skipped, "standard_source_rows": len(standard_rows), "standard_inserted": standard_inserted, "standard_skipped": standard_skipped, } def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--anomaly-file", type=Path, required=True) parser.add_argument("--standard-file", type=Path, required=True) args = parser.parse_args() logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s %(message)s") if not args.anomaly_file.is_file(): parser.error(f"anomaly file does not exist: {args.anomaly_file}") if not args.standard_file.is_file(): parser.error(f"standard file does not exist: {args.standard_file}") result = asyncio.run(import_workbooks(args.anomaly_file, args.standard_file)) for key, value in result.items(): print(f"{key}={value}") if __name__ == "__main__": main()