PayrollManagementSystem/backend/app/services/tplus/reader.py

"""T+ 六类源表的只读读取器。"""

from collections.abc import Iterator
from datetime import datetime
from pathlib import Path
from typing import Any

from sqlalchemy import BigInteger, Integer, bindparam, text
from sqlalchemy.orm import Session

from app.services.tplus.mapping import TPLUS_MAPPINGS, get_mapping

_ENTITY_QUERY_FILES = {
    "personnel": "personnel.sql",
    "products": "products.sql",
    "customers": "customers.sql",
    "businesses": "businesses.sql",
    "shipments": "shipments.sql",
    "receipts": "receipts.sql",
    "business_items": "business_items.sql",
    "shipment_items": "shipment_items.sql",
    "receipt_items": "receipt_items.sql",
}


class TPlusReader:
    """执行白名单 SQL 文件，禁止调用方传入任意表名或 SQL。"""

    def __init__(self, session: Session, query_dir: Path | None = None) -> None:
        self.session = session
        self.query_dir = query_dir or self._default_query_dir()

    @staticmethod
    def _default_query_dir() -> Path:
        return Path(__file__).resolve().parents[4] / "db" / "queries" / "tplus"

    def count_entities(self, entity_types) -> dict[str, int]:
        """只统计选定的白名单范围，保留读取 SQL 的同一源过滤语义。"""
        entities = tuple(dict.fromkeys(entity_types))
        for entity in entities:
            get_mapping(entity)
        if not entities:
            return {}
        statement = text((self.query_dir / "progress" / "counts.sql").read_text(encoding="utf-8")).bindparams(
            bindparam("entities", expanding=True)
        )
        return {row.entity_type: int(row.total_records) for row in self.session.execute(
            statement, {"entities": entities}
        )}

    def iter_entity(
        self,
        entity_type: str,
        *,
        watermark: datetime | None = None,
        batch_size: int = 500,
    ) -> Iterator[dict[str, Any]]:
        """按实体流式读取，水位只通过绑定参数传入。"""
        if entity_type not in _ENTITY_QUERY_FILES:
            raise KeyError(f"不支持的 T+ 实体类型：{entity_type}")
        if batch_size < 1 or batch_size > 5000:
            raise ValueError("batch_size 必须在 1 到 5000 之间")
        if watermark is not None:
            raise ValueError("增量水位语义尚未确认，当前只支持全量读取和哈希去重")
        get_mapping(entity_type)
        sql_path = self.query_dir / _ENTITY_QUERY_FILES[entity_type]
        if not sql_path.is_file():
            raise FileNotFoundError(f"缺少 T+ 查询文件：{sql_path}")

        statement = text(sql_path.read_text(encoding="utf-8")).bindparams(
            bindparam("batch_size", type_=Integer()), bindparam("after_id", type_=BigInteger())
        )
        after_id = -2147483649  # 所有源表主键均已核验为 SQL Server int，包含其最小值。
        while True:
            result = self.session.execute(
                statement, {"batch_size": batch_size, "after_id": after_id}
            )
            try:
                rows = [dict(row) for row in result.mappings().all()]
            finally:
                result.close()
            if not rows:
                return
            next_id = int(rows[-1]["id"])
            if next_id <= after_id:
                raise ValueError("T+ 源主键未递增，停止读取以避免重复循环")
            yield from rows
            after_id = next_id

    def fetch_entity(
        self,
        entity_type: str,
        *,
        watermark: datetime | None = None,
        max_rows: int = 5000,
    ) -> list[dict[str, Any]]:
        """读取有限数量的实体行，供核验命令和单元测试使用。"""
        if max_rows < 1:
            raise ValueError("max_rows 必须大于 0")
        rows: list[dict[str, Any]] = []
        for row in self.iter_entity(entity_type, watermark=watermark):
            rows.append(row)
            if len(rows) >= max_rows:
                break
        return rows

    def inventory(self) -> tuple[str, ...]:
        """返回当前适配器支持的实体类型。"""
        return tuple(mapping.entity_type for mapping in TPLUS_MAPPINGS)