#!/usr/bin/env python3 from __future__ import annotations import argparse import json import os import sys import uuid from dataclasses import asdict from datetime import UTC, datetime from pathlib import Path from typing import Any from sqlalchemy import create_engine, text from sqlalchemy.engine import Connection from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import Session from sqlalchemy.pool import NullPool SERVER_DIR = Path(__file__).resolve().parents[1] SRC_DIR = SERVER_DIR / "src" if str(SRC_DIR) not in sys.path: sys.path.insert(0, str(SRC_DIR)) from app.db.maintenance_database_target import ( # noqa: E402 MaintenanceDatabaseTarget, MaintenanceDatabaseTargetError, validate_maintenance_database_target, ) from app.db.migration_preflight import ( # noqa: E402 MigrationPreflightError, validate_migration_state, ) from app.services.expense_case_legacy_backfill import ( # noqa: E402 DEFAULT_BATCH_SIZE, MAX_BATCH_SIZE, ExpenseCaseLegacyBackfillService, LegacyBackfillCursor, ) REQUIRED_ALEMBIC_REVISION = "20260713_0002" EXIT_CONFIGURATION = 2 EXIT_SAFETY = 3 EXIT_LOCKED = 4 EXIT_CONFLICT = 5 EXIT_RUNTIME = 6 class BackfillCommandError(RuntimeError): def __init__( self, message: str, *, exit_code: int, code: str, details: dict[str, Any] | None = None, ) -> None: super().__init__(message) self.exit_code = exit_code self.code = code self.details = details def parse_created_before(value: str) -> datetime: normalized = str(value or "").strip() if normalized.endswith("Z"): normalized = f"{normalized[:-1]}+00:00" try: parsed = datetime.fromisoformat(normalized) except ValueError as exc: raise argparse.ArgumentTypeError( "--created-before 必须是带时区的 ISO 8601 时间,例如 2026-07-14T00:00:00Z" ) from exc if parsed.tzinfo is None or parsed.utcoffset() is None: raise argparse.ArgumentTypeError("--created-before 必须显式包含时区") return parsed.astimezone(UTC) def positive_int(value: str) -> int: try: parsed = int(value) except ValueError as exc: raise argparse.ArgumentTypeError("必须是正整数") from exc if parsed < 1: raise argparse.ArgumentTypeError("必须是正整数") return parsed def non_empty_text(value: str) -> str: normalized = str(value or "").strip() if not normalized: raise argparse.ArgumentTypeError("不能为空") return normalized def batch_size(value: str) -> int: parsed = positive_int(value) if parsed > MAX_BATCH_SIZE: raise argparse.ArgumentTypeError(f"不能超过 {MAX_BATCH_SIZE}") return parsed def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description="把迁移前 ExpenseClaim 诚实地接入统一费用事件;默认只预览。", ) mode = parser.add_mutually_exclusive_group() mode.add_argument("--dry-run", action="store_true", help="只读预览(默认)。") mode.add_argument("--apply", action="store_true", help="按批写入历史快照事件。") parser.add_argument( "--tenant-id", required=True, type=non_empty_text, help="旧单明确归属的租户 ID。", ) parser.add_argument( "--created-before", required=True, type=parse_created_before, help="仅处理该时刻之前创建的单据,必须包含时区。", ) parser.add_argument("--batch-size", type=batch_size, default=DEFAULT_BATCH_SIZE) parser.add_argument( "--max-claims", type=positive_int, help="最多扫描的历史单据数,用于 canary。", ) parser.add_argument("--sample-limit", type=positive_int, default=20) parser.add_argument("--expected-host", required=True) parser.add_argument("--expected-database", required=True) parser.add_argument( "--confirm-target", help="apply 时必须精确等于解析后的 host:port/database。", ) parser.add_argument( "--allow-non-disposable-target", action="store_true", help="允许在非 probe 数据库 apply;仍需精确确认目标。", ) return parser def _cursor_payload(cursor: LegacyBackfillCursor | None) -> dict[str, str] | None: if cursor is None: return None return { "created_at": cursor.created_at.astimezone(UTC).isoformat().replace("+00:00", "Z"), "claim_id": cursor.claim_id, } def _item_payload(item: Any) -> dict[str, str]: payload = asdict(item) payload["disposition"] = item.disposition.value return payload def _base_summary( *, mode: str, target: MaintenanceDatabaseTarget, tenant_id: str, created_before: datetime, revision: str, ) -> dict[str, Any]: return { "mode": mode, "database": { "target": target.exact_target, "url": target.sanitized_url, "disposable": target.is_disposable, "revision": revision, }, "tenant_id": tenant_id, "created_before": created_before.isoformat().replace("+00:00", "Z"), "inspected": 0, "eligible": 0, "already_linked": 0, "conflicts": 0, "created": 0, "batches": 0, "limited": False, "last_cursor": None, "samples": [], } def _page_limit(*, configured: int, remaining: int | None) -> int: return configured if remaining is None else min(configured, remaining) def _execution_progress(summary: dict[str, Any], **extra: Any) -> dict[str, Any]: progress = { "run_id": summary.get("run_id"), "partial_commit": bool(summary.get("batches")), "committed_batches": int(summary.get("batches") or 0), "committed_claims": int(summary.get("inspected") or 0), "created": int(summary.get("created") or 0), "already_linked": int(summary.get("already_linked") or 0), "last_committed_cursor": summary.get("last_cursor"), } progress.update(extra) return progress def _preview( session: Session, *, tenant_id: str, created_before: datetime, configured_batch_size: int, max_claims: int | None, sample_limit: int, summary: dict[str, Any], ) -> None: service = ExpenseCaseLegacyBackfillService( session, tenant_id=tenant_id, cutoff=created_before, ) cursor: LegacyBackfillCursor | None = None remaining = max_claims while remaining is None or remaining > 0: page = service.preview( batch_size=_page_limit(configured=configured_batch_size, remaining=remaining), after=cursor, ) if not page.items: break summary["batches"] += 1 summary["inspected"] += page.inspected summary["eligible"] += page.eligible summary["already_linked"] += page.linked summary["conflicts"] += page.conflicts available_samples = max(0, sample_limit - len(summary["samples"])) summary["samples"].extend(_item_payload(item) for item in page.items[:available_samples]) cursor = page.next_cursor summary["last_cursor"] = _cursor_payload(cursor) if remaining is not None: remaining -= page.inspected if not page.has_more: break summary["limited"] = bool(remaining == 0 and page.has_more) if "page" in locals() else False def _acquire_advisory_lock(connection: Connection, tenant_id: str) -> str: lock_name = f"legacy-expense-case-backfill:{tenant_id}" acquired = connection.scalar( text("SELECT pg_try_advisory_lock(hashtextextended(:lock_name, 0))"), {"lock_name": lock_name}, ) connection.commit() if not acquired: raise BackfillCommandError( "同一租户已有历史费用事件回填任务正在运行。", exit_code=EXIT_LOCKED, code="advisory_lock_unavailable", ) return lock_name def _release_advisory_lock(connection: Connection, lock_name: str) -> None: if connection.in_transaction(): connection.rollback() connection.execute( text("SELECT pg_advisory_unlock(hashtextextended(:lock_name, 0))"), {"lock_name": lock_name}, ) connection.commit() def _apply( connection: Connection, *, tenant_id: str, created_before: datetime, configured_batch_size: int, max_claims: int | None, sample_limit: int, summary: dict[str, Any], ) -> None: lock_name = _acquire_advisory_lock(connection, tenant_id) run_id = f"historical-import-{uuid.uuid4().hex}" backfilled_at = datetime.now(UTC) summary["run_id"] = run_id summary["backfilled_at"] = backfilled_at.isoformat().replace("+00:00", "Z") cursor: LegacyBackfillCursor | None = None remaining = max_claims try: with Session(bind=connection, autoflush=False, expire_on_commit=False) as session: service = ExpenseCaseLegacyBackfillService( session, tenant_id=tenant_id, cutoff=created_before, ) while remaining is None or remaining > 0: session.execute(text("SET LOCAL lock_timeout = '5s'")) result = service.apply_batch( run_id=run_id, batch_size=_page_limit( configured=configured_batch_size, remaining=remaining, ), after=cursor, backfilled_at=backfilled_at, ) if not result.items: session.rollback() break if result.conflicts: session.rollback() conflict_ids = [ item.claim_id for item in result.items if item.disposition.value == "conflict" ] raise BackfillCommandError( f"检测到 {result.conflicts} 个数据冲突,当前批次已回滚:" f"{', '.join(conflict_ids)}", exit_code=EXIT_CONFLICT, code="legacy_claim_conflict", details=_execution_progress( summary, current_batch_conflict_ids=conflict_ids, ), ) session.commit() summary["batches"] += 1 summary["inspected"] += result.inspected summary["created"] += result.created summary["already_linked"] += result.skipped_linked available_samples = max(0, sample_limit - len(summary["samples"])) summary["samples"].extend( _item_payload(item) for item in result.items[:available_samples] ) cursor = result.next_cursor summary["last_cursor"] = _cursor_payload(cursor) if remaining is not None: remaining -= result.inspected if not result.has_more: break summary["limited"] = bool(remaining == 0 and "result" in locals() and result.has_more) finally: _release_advisory_lock(connection, lock_name) def _verify_connected_database( connection: Connection, *, expected_database: str, ) -> tuple[str, str]: database_name, database_user = connection.execute( text("SELECT current_database(), current_user") ).one() if str(database_name) != expected_database: raise BackfillCommandError( "连接后的 current_database() 与 DATABASE_URL 不一致。", exit_code=EXIT_SAFETY, code="connected_database_mismatch", ) return str(database_name), str(database_user) def run(args: argparse.Namespace) -> dict[str, Any]: database_url = os.environ.get("DATABASE_URL", "") if args.apply and str(args.confirm_target or "").strip() == "": raise BackfillCommandError( "--apply 必须提供 --confirm-target 精确确认数据库目标。", exit_code=EXIT_SAFETY, code="confirm_target_required", ) target = validate_maintenance_database_target( database_url, expected_host=args.expected_host, expected_database=args.expected_database, apply=args.apply, allow_non_disposable=args.allow_non_disposable_target, confirm_target=args.confirm_target, ) engine = create_engine( database_url, pool_pre_ping=True, poolclass=NullPool, ) try: with engine.connect() as connection: connected_database, connected_user = _verify_connected_database( connection, expected_database=target.database, ) state = validate_migration_state(connection) if state.revision != REQUIRED_ALEMBIC_REVISION: raise BackfillCommandError( "数据库必须先升级到费用事件迁移 head:" f"actual={state.revision or 'unversioned/base'}; " f"required={REQUIRED_ALEMBIC_REVISION}", exit_code=EXIT_SAFETY, code="migration_revision_mismatch", ) connection.rollback() summary = _base_summary( mode="apply" if args.apply else "dry-run", target=target, tenant_id=args.tenant_id, created_before=args.created_before, revision=state.revision, ) summary["database"]["connected_database"] = connected_database summary["database"]["connected_user"] = connected_user with Session(bind=connection, autoflush=False) as preview_session: _preview( preview_session, tenant_id=args.tenant_id, created_before=args.created_before, configured_batch_size=args.batch_size, max_claims=args.max_claims, sample_limit=args.sample_limit, summary=summary, ) preview_session.rollback() if not args.apply: summary["would_create_cases"] = summary["eligible"] summary["would_create_links"] = summary["eligible"] summary["would_create_events"] = summary["eligible"] return summary if summary["conflicts"]: raise BackfillCommandError( f"预览发现 {summary['conflicts']} 个数据冲突;未执行任何写入。", exit_code=EXIT_CONFLICT, code="legacy_claim_conflict", ) preview = { "inspected": summary["inspected"], "eligible": summary["eligible"], "already_linked": summary["already_linked"], "samples": summary["samples"], } summary.update( inspected=0, eligible=preview["eligible"], already_linked=0, conflicts=0, created=0, batches=0, limited=False, last_cursor=None, samples=[], preview=preview, ) try: _apply( connection, tenant_id=args.tenant_id, created_before=args.created_before, configured_batch_size=args.batch_size, max_claims=args.max_claims, sample_limit=args.sample_limit, summary=summary, ) except BackfillCommandError: raise except SQLAlchemyError as exc: raise BackfillCommandError( "数据库执行失败;已提交批次不会回滚,请依据进度摘要安全重跑。", exit_code=EXIT_RUNTIME, code="database_runtime_error", details=_execution_progress(summary), ) from exc return summary finally: engine.dispose() def _error_payload( exc: Exception, *, code: str, details: dict[str, Any] | None = None, ) -> dict[str, Any]: payload: dict[str, Any] = {"status": "error", "code": code, "message": str(exc)} if details is not None: payload["details"] = details return payload def main(argv: list[str] | None = None) -> int: args = build_parser().parse_args(argv) try: payload = run(args) except BackfillCommandError as exc: print( json.dumps( _error_payload(exc, code=exc.code, details=exc.details), ensure_ascii=False, ), file=sys.stderr, ) return exc.exit_code except MaintenanceDatabaseTargetError as exc: print( json.dumps(_error_payload(exc, code=exc.code), ensure_ascii=False), file=sys.stderr, ) return EXIT_SAFETY except MigrationPreflightError as exc: print( json.dumps(_error_payload(exc, code="migration_preflight_failed"), ensure_ascii=False), file=sys.stderr, ) return EXIT_SAFETY except (SQLAlchemyError, OSError) as exc: print( json.dumps(_error_payload(exc, code="database_runtime_error"), ensure_ascii=False), file=sys.stderr, ) return EXIT_RUNTIME print(json.dumps(payload, ensure_ascii=False, indent=2, default=str)) return 0 if __name__ == "__main__": raise SystemExit(main())