feat(expenses): backfill historical claims into expense cases
This commit is contained in:
caoxiaozhu
532
server/scripts/backfill_legacy_expense_claim_cases.py
Normal file
532
server/scripts/backfill_legacy_expense_claim_cases.py
Normal file
@@ -0,0 +1,532 @@
|
||||
#!/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())
|
||||
Reference in New Issue
Block a user