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

"""周期检查只排队源采集：稳定意图键、有限重试、独立心跳和只读异常概况。"""

import hashlib
import logging
import threading
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from pathlib import Path
from uuid import NAMESPACE_URL, uuid4, uuid5

from sqlalchemy import func, inspect, select, text
from sqlalchemy.dialects.mysql import insert
from sqlalchemy.orm import sessionmaker

from app.core.config import get_settings
from app.core.tplus_db import resolve_tplus_settings
from app.models import TPlusAutomationState, TPlusSyncBatch, TPlusSyncError, TPlusWorkRun
from app.services.tplus import task_service
from app.services.tplus.errors import TRANSIENT_CODES
from app.services.tplus.mapping import DETAIL_MAPPINGS, TPLUS_MAPPINGS

logger = logging.getLogger(__name__)
SHANGHAI = timezone(timedelta(hours=8))
ENTITIES = tuple(sorted(m.entity_type for m in (*TPLUS_MAPPINGS, *DETAIL_MAPPINGS)))


def local_now():
    return datetime.now(SHANGHAI).replace(tzinfo=None)


def source_identity(settings):
    source = resolve_tplus_settings(settings)
    value = f"{source.t_plus_host.casefold()}:{source.t_plus_port}/{source.t_plus_database}"
    return hashlib.sha256(value.encode()).hexdigest(), source.t_plus_database


def due_period(now, weekday=0, clock="02:00"):
    """取最近一次已到期时段；只补这一个周期，不逐周补发历史任务。"""
    if now.tzinfo:
        now = now.astimezone(SHANGHAI).replace(tzinfo=None)
    hour, minute = map(int, clock.split(":"))
    day = now.date() - timedelta(days=(now.weekday() - weekday) % 7)
    due = datetime.combine(day, datetime.min.time()).replace(hour=hour, minute=minute)
    if due > now:
        due -= timedelta(days=7)
    year, week, _ = due.isocalendar()
    return f"{year}-W{week:02d}", due


def intent_key(source_key, period, attempt=0):
    return str(uuid5(NAMESPACE_URL, f"payroll:tplus:weekly:{source_key}:{period}:{attempt}"))


def _state(session, key, database, *, lock=False):
    if lock and session.bind.dialect.name == "mysql":
        stmt = insert(TPlusAutomationState).values(source_key=key, source_database=database)
        session.execute(stmt.on_duplicate_key_update(source_key=stmt.inserted.source_key))
    statement = select(TPlusAutomationState).where(TPlusAutomationState.source_key == key)
    if lock:
        statement = statement.with_for_update().execution_options(populate_existing=True)
    row = session.scalar(statement)
    if row is None and lock:
        row = TPlusAutomationState(source_key=key, source_database=database, retry_attempts=0, worker_phase="STOPPED")
        session.add(row)
        session.flush()
    return row


@contextmanager
def maintenance_lock(session, key):
    """维护检查独立于长采集锁；同一源的两个检查不能重复创建重试。"""
    if session.bind.dialect.name != "mysql":
        yield True
        return
    connection = session.get_bind().connect()
    name = "ys:tplus:maintain:" + key[:32]
    acquired = False
    try:
        acquired = connection.execute(text("SELECT GET_LOCK(:name,0)"), {"name": name}).scalar() == 1
        yield acquired
    finally:
        if acquired:
            try:
                if connection.execute(text("SELECT RELEASE_LOCK(:name)"), {"name": name}).scalar() != 1:
                    connection.invalidate()
            except Exception:
                connection.invalidate()
        connection.close()


def _has_active(session):
    # 实际读取仍受原worker命名锁保护；排队检查防止周期任务积压。
    for model in (TPlusSyncBatch, TPlusWorkRun):
        if session.scalar(select(model.id).where(model.status.in_(("QUEUED", "RUNNING"))).limit(1)) is not None:
            return True
    return False


def _enqueue(session, state, period, attempt, parent=None):
    values = {
        "batch_no": task_service.new_batch_no(), "request_key": intent_key(state.source_key, period, attempt),
        "sync_mode": "FULL", "status": "QUEUED", "source_database": state.source_database,
        "entity_types_json": list(ENTITIES), "parent_batch_no": parent,
        "trigger_type": "AUTO_RETRY" if attempt else "SCHEDULED", "automatic_attempt": attempt,
        "source_record_count": 0, "staged_record_count": 0, "failed_record_count": 0, "skipped_record_count": 0,
    }
    batch, inserted = task_service._insert(session, values)
    if inserted and parent:
        from app.repositories.tplus_repository import increment_retry_errors

        increment_retry_errors(session, parent)
    state.scheduled_period = period
    state.current_batch_no = batch.batch_no
    state.retry_attempts = attempt
    state.next_retry_at = None
    if not attempt:
        state.root_batch_no = batch.batch_no
    return batch


def _cycle_current(session, state):
    """锁住本周期重试树；人工恢复也算本周期进度，自动次数不能被重置。"""
    if not state.root_batch_no:
        return None
    root = session.scalar(select(TPlusSyncBatch).where(
        TPlusSyncBatch.batch_no == state.root_batch_no,
    ).with_for_update().execution_options(populate_existing=True))
    if root is None:
        return None
    current, frontier, seen = root, [root.batch_no], {root.batch_no}
    while frontier:
        children = list(session.scalars(select(TPlusSyncBatch).where(
            TPlusSyncBatch.parent_batch_no.in_(frontier),
            TPlusSyncBatch.source_database == state.source_database,
        ).order_by(TPlusSyncBatch.id).with_for_update().execution_options(populate_existing=True)))
        frontier = []
        for child in children:
            if child.batch_no in seen:
                continue
            seen.add(child.batch_no)
            frontier.append(child.batch_no)
            state.retry_attempts = max(state.retry_attempts or 0, child.automatic_attempt or 0)
            if child.id > current.id:
                current = child
    state.current_batch_no = current.batch_no
    return current


def tick(session, *, settings=None, now=None):
    settings, now = settings or get_settings(), now or local_now()
    from app.services.tplus.automation_config import configuration

    if not configuration(session, settings)["enabled"]:
        return {"status": "DISABLED", "queued": False}
    if settings.t_plus_stop_file and Path(settings.t_plus_stop_file).exists():
        return {"status": "STOP_REQUESTED", "queued": False}
    key, database = source_identity(settings)
    with maintenance_lock(session, key) as acquired:
        if not acquired:
            return {"status": "BUSY", "queued": False}
        session.rollback()
        state = _state(session, key, database, lock=True)
        state.maintenance_at = now
        period, _ = due_period(now, settings.t_plus_schedule_weekday, settings.t_plus_schedule_time)
        current = _cycle_current(session, state)
        if current and current.status == "SUCCEEDED":
            state.last_success_at = current.finished_at
            state.next_retry_at = None
        if _has_active(session):
            session.commit()
            return {"status": "DEFERRED", "queued": False, "period": period}
        batch = None
        if state.scheduled_period != period:
            batch = _enqueue(session, state, period, 0)
        elif current and current.status == "FAILED" and (state.retry_attempts or 0) < 2:
            codes = set(session.scalars(select(TPlusSyncError.error_code).where(TPlusSyncError.batch_no == current.batch_no)))
            if codes and codes <= TRANSIENT_CODES and current.finished_at is not None:
                delay = (5, 15)[state.retry_attempts or 0]
                state.next_retry_at = current.finished_at + timedelta(minutes=delay)
                if now >= state.next_retry_at:
                    batch = _enqueue(session, state, period, (state.retry_attempts or 0) + 1, current.batch_no)
            else:
                state.next_retry_at = None
        elif current and current.status in ("FAILED", "PARTIAL"):
            state.next_retry_at = None
        session.commit()
        return {"status": "QUEUED" if batch else "CHECKED", "queued": batch is not None,
            "batch_no": batch.batch_no if batch else None, "period": period}


def heartbeat(session_factory, settings, owner, phase, *, now=None, claim=False):
    key, database = source_identity(settings)
    with session_factory() as session:
        state = _state(session, key, database, lock=True)
        if state.worker_owner not in (None, owner) and not claim:
            session.rollback()
            return
        state.worker_owner, state.worker_phase = owner, phase
        state.heartbeat_at = now or local_now()
        session.commit()


class WorkerHeartbeat:
    """独立会话和线程写心跳，长查询不阻塞；只在worker已持采集锁时启动。"""
    def __init__(self, session, phase="RUNNING", interval=30, owner=None):
        self.factory = sessionmaker(session.get_bind(), expire_on_commit=False)
        self.settings = get_settings()
        self.owner, self.phase, self.interval = owner or str(uuid4()), phase, interval
        self.stop = threading.Event()
        self.thread = None

    def __enter__(self):
        try:
            source_identity(self.settings)
            with self.factory() as session:
                if not inspect(session.connection()).has_table(TPlusAutomationState.__tablename__):
                    return self
            heartbeat(self.factory, self.settings, self.owner, self.phase, claim=True)
        except Exception:
            logger.warning("T+心跳暂不可写，请检查迁移与运行配置")
            return self
        self.thread = threading.Thread(target=self._pulse, name="tplus-heartbeat", daemon=True)
        self.thread.start()
        return self

    def _pulse(self):
        while not self.stop.wait(self.interval):
            try:
                heartbeat(self.factory, self.settings, self.owner, self.phase)
            except Exception:
                logger.warning("T+进程心跳更新失败")

    def __exit__(self, *_args):
        self.stop.set()
        if self.thread:
            self.thread.join(timeout=5)
            try:
                heartbeat(self.factory, self.settings, self.owner, "IDLE")
            except Exception:
                logger.warning("T+进程状态更新失败")


def overview(session, *, settings=None, now=None):
    settings, now = settings or get_settings(), now or local_now()
    base = {"enabled": settings.t_plus_automation_enabled, "weekday": settings.t_plus_schedule_weekday,
        "time": settings.t_plus_schedule_time, "timezone": "Asia/Shanghai", "worker_phase": "UNKNOWN",
        "heartbeat_at": None, "last_success_at": None, "next_retry_at": None, "alerts": []}
    from app.services.tplus.automation_config import configuration

    base.update(configuration(session, settings))
    base["stop_requested"] = bool(settings.t_plus_stop_file and Path(settings.t_plus_stop_file).exists())
    try:
        key, database = source_identity(settings)
    except RuntimeError:
        database = task_service.history_source_database(session)
        if database:
            base["last_success_at"] = session.scalar(
                select(func.max(TPlusSyncBatch.finished_at)).where(
                    TPlusSyncBatch.source_database == database,
                    TPlusSyncBatch.status == "SUCCEEDED",
                    TPlusSyncBatch.sync_mode == "FULL",
                )
            )
        return base
    if not inspect(session.connection()).has_table(TPlusAutomationState.__tablename__):
        base["alerts"] = [{"code": "MIGRATION_REQUIRED", "message": "运行状态结构尚未升级"}]
        return base
    state = _state(session, key, database)
    if state:
        base.update(worker_phase=state.worker_phase, heartbeat_at=state.heartbeat_at,
            last_success_at=state.last_success_at, next_retry_at=state.next_retry_at)
    alerts = base["alerts"]
    latest = session.scalar(select(TPlusSyncBatch).where(TPlusSyncBatch.source_database == database).order_by(TPlusSyncBatch.id.desc()).limit(1))
    if latest and latest.status in ("FAILED", "PARTIAL"):
        alerts.append({"code": "SYNC_REQUIRES_ATTENTION", "message": "最近采集失败或部分失败，请核对后处理", "batch_no": latest.batch_no})
    queued = session.scalar(select(TPlusSyncBatch).where(TPlusSyncBatch.source_database == database,
        TPlusSyncBatch.status == "QUEUED", TPlusSyncBatch.create_time <= now - timedelta(minutes=15)).order_by(TPlusSyncBatch.id).limit(1))
    if queued:
        alerts.append({"code": "QUEUE_OVERDUE", "message": "采集任务排队超过15分钟", "batch_no": queued.batch_no})
    last_success = session.scalar(select(func.max(TPlusSyncBatch.finished_at)).where(TPlusSyncBatch.source_database == database, TPlusSyncBatch.status == "SUCCEEDED", TPlusSyncBatch.sync_mode == "FULL"))
    base["last_success_at"] = last_success
    if base["enabled"] and not base["stop_requested"]:
        if last_success is None or now - last_success > timedelta(days=8):
            alerts.append({"code": "SYNC_STALE", "message": "尚无成功采集或超过8天未成功同步"})
        if state is None or state.heartbeat_at is None or now - state.heartbeat_at > timedelta(minutes=2) or state.worker_phase == "STOPPED":
            alerts.append({"code": "WORKER_OFFLINE", "message": "采集进程心跳中断，请检查托管任务"})
    return base