PayrollManagementSystem/backend/app/repositories/tplus_repository.py

"""T+ 管理库访问；只 flush，事务由 service 管理。"""

from sqlalchemy import case, func, select, update
from sqlalchemy.orm import Session, defer

from app.models import TPlusSyncBatch, TPlusSyncError


def get_batch(
    session: Session, batch_no: str, *, for_update: bool = False, with_result: bool = False
):
    """按批次号查询；普通查询不加载核验大字段。"""
    statement = select(TPlusSyncBatch).where(TPlusSyncBatch.batch_no == batch_no)
    if not with_result:
        statement = statement.options(defer(TPlusSyncBatch.audit_result_json))
    if for_update:
        statement = statement.with_for_update().execution_options(populate_existing=True)
    return session.scalar(statement)


def get_request(session: Session, key: str, *, for_update: bool = False):
    """锁读用于唯一约束竞争后的当前读，避免旧事务快照看不到新任务。"""
    statement = select(TPlusSyncBatch).where(TPlusSyncBatch.request_key == key)
    if for_update:
        statement = statement.with_for_update().execution_options(populate_existing=True)
    return session.scalar(statement)


def active_retry(session: Session, parent_no: str):
    """调用方先锁父批次，再用当前读判断是否已有未完成重试。"""
    return session.scalar(
        select(TPlusSyncBatch)
        .where(
            TPlusSyncBatch.parent_batch_no == parent_no,
            TPlusSyncBatch.status.in_(("QUEUED", "RUNNING")),
        )
        .with_for_update()
    )


def increment_retry_errors(session: Session, parent_no: str) -> None:
    """创建新重试时对原异常累计重试次数。"""
    session.execute(
        update(TPlusSyncError)
        .where(TPlusSyncError.batch_no == parent_no)
        .values(retry_count=TPlusSyncError.retry_count + 1)
    )
    session.flush()


def next_queued(session: Session):
    """调用方持有账套执行锁后领取最早排队任务。"""
    return session.scalar(
        select(TPlusSyncBatch)
        .where(TPlusSyncBatch.status == "QUEUED")
        .order_by(TPlusSyncBatch.id)
        .limit(1)
        .with_for_update()
    )


def running_batches(session: Session) -> list[TPlusSyncBatch]:
    """只在持有执行锁时使用，不能依据时间戳判断任务已经停止。"""
    return list(
        session.scalars(
            select(TPlusSyncBatch).where(TPlusSyncBatch.status == "RUNNING").with_for_update()
        ).all()
    )


def list_batches(session: Session, query):
    """分页只读元数据，核验报告和原异常不随列表加载。"""
    filters = []
    if query.status:
        filters.append(TPlusSyncBatch.status == query.status)
    if query.task_type:
        filters.append(
            TPlusSyncBatch.sync_mode == ("AUDIT" if query.task_type == "AUDIT" else "FULL")
        )
    total = session.scalar(select(func.count()).select_from(TPlusSyncBatch).where(*filters))
    rows = session.scalars(
        select(TPlusSyncBatch)
        .where(*filters)
        .options(defer(TPlusSyncBatch.audit_result_json), defer(TPlusSyncBatch.error_message))
        .order_by(TPlusSyncBatch.id.desc())
        .offset(query.offset)
        .limit(query.page_size)
    ).all()
    return rows, total or 0


def active_retry_parents(session: Session, numbers: list[str]) -> set[str]:
    """整页一次获取未完成的重试，避免逐行查询。"""
    if not numbers:
        return set()
    return set(
        session.scalars(
            select(TPlusSyncBatch.parent_batch_no).where(
                TPlusSyncBatch.parent_batch_no.in_(numbers),
                TPlusSyncBatch.status.in_(("QUEUED", "RUNNING")),
            )
        ).all()
    )


def list_errors(session: Session, number: str, query):
    """只读取当前批次的异常定位元数据。"""
    filters = [TPlusSyncError.batch_no == number]
    if query.entity_type:
        filters.append(TPlusSyncError.entity_type == query.entity_type)
    total = session.scalar(select(func.count()).select_from(TPlusSyncError).where(*filters))
    rows = session.scalars(
        select(TPlusSyncError)
        .where(*filters)
        .options(defer(TPlusSyncError.error_message))
        .order_by(TPlusSyncError.id.desc())
        .offset(query.offset)
        .limit(query.page_size)
    ).all()
    return rows, total or 0


def status_counts(session: Session) -> dict[str, int]:
    """汇总任务表，不扫描原始快照。"""
    return dict(
        session.execute(
            select(TPlusSyncBatch.status, func.count()).group_by(TPlusSyncBatch.status)
        ).all()
    )


def progress_reports(session: Session, numbers):
    if not numbers:
        return {}
    rows = session.execute(select(TPlusSyncBatch.batch_no,
        TPlusSyncBatch.audit_result_json["progress"],
        TPlusSyncBatch.audit_result_json["observed_entities"],
    ).where(TPlusSyncBatch.batch_no.in_(numbers)))
    return {number: {"progress": progress, "observed_entities": observed or []}
            for number, progress, observed in rows}


def current_batch(session: Session):
    return session.scalar(select(TPlusSyncBatch).order_by(
        case((TPlusSyncBatch.status == "RUNNING", 0), (TPlusSyncBatch.status == "QUEUED", 1), else_=2),
        TPlusSyncBatch.id.desc(),
    ).limit(1))


def latest_audit(session: Session):
    """最近完成的核验；返回批次供查询层过滤报告。"""
    return session.scalar(
        select(TPlusSyncBatch)
        .where(
            TPlusSyncBatch.sync_mode == "AUDIT",
            TPlusSyncBatch.finished_at.is_not(None),
        )
        .order_by(TPlusSyncBatch.id.desc())
        .limit(1)
    )


def total_staged_records(session: Session) -> int:
    """从批次元数据汇总累计新增数，避免扫描原始快照。"""
    return (
        session.scalar(select(func.coalesce(func.sum(TPlusSyncBatch.staged_record_count), 0))) or 0
    )