Files
X-Financial/server/scripts/backfill_legacy_expense_claim_cases.py

533 lines
18 KiB
Python
Raw Normal View History

#!/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())