97 lines
3.9 KiB
Python
97 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import sys
|
|
from decimal import Decimal
|
|
|
|
from sqlalchemy import text
|
|
|
|
from app.enums.credit_balance import CreditBalanceSourceType, CreditLevel
|
|
from app.enums.credit_record import CreditRecordType
|
|
from app.models.base import async_session
|
|
from app.services.credit.ledger_service import grant_credits
|
|
from app.services.credit.time_policy import add_natural_months
|
|
from app.services.credit.utils import to_credit_decimal, utc_now
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description="将 users.credits 一次性迁移到动态积分余额表。")
|
|
parser.add_argument("--batch-size", type=int, default=500)
|
|
parser.add_argument("--dry-run", action="store_true")
|
|
parser.add_argument("--commit", action="store_true")
|
|
parser.add_argument("--after-id", default="")
|
|
return parser
|
|
|
|
|
|
async def amain(argv: list[str]) -> int:
|
|
args = build_parser().parse_args(argv)
|
|
if args.dry_run and args.commit:
|
|
raise SystemExit("--dry-run 和 --commit 不能同时使用")
|
|
do_commit = bool(args.commit)
|
|
batch_size = max(1, min(int(args.batch_size or 500), 5000))
|
|
migration_time = utc_now()
|
|
cursor = str(args.after_id or "")
|
|
stats = {"positive_users": 0, "zero_users": 0, "negative_reset_users": 0, "migrated_credits": "0.00", "last_id": cursor}
|
|
total = Decimal("0.00")
|
|
|
|
while True:
|
|
async with async_session() as db:
|
|
rows = (await db.execute(
|
|
text("SELECT id, credits FROM users WHERE id > :cursor ORDER BY id ASC LIMIT :limit"),
|
|
{"cursor": cursor, "limit": batch_size},
|
|
)).mappings().all()
|
|
if not rows:
|
|
break
|
|
try:
|
|
for row in rows:
|
|
user_id = str(row["id"])
|
|
legacy = to_credit_decimal(row.get("credits") or 0)
|
|
cursor = user_id
|
|
stats["last_id"] = cursor
|
|
if legacy > 0:
|
|
stats["positive_users"] += 1
|
|
total += legacy
|
|
if do_commit:
|
|
await grant_credits(
|
|
db,
|
|
user_id=user_id,
|
|
amount=legacy,
|
|
description="历史用户积分迁移",
|
|
source_type=CreditBalanceSourceType.LEGACY_MIGRATION.value,
|
|
valid_from=migration_time,
|
|
expires_at=add_natural_months(migration_time, 1),
|
|
credit_level=CreditLevel.GENERAL.value,
|
|
source_id=user_id,
|
|
related_id=user_id,
|
|
record_type=CreditRecordType.RECHARGE.value,
|
|
biz_key=f"legacy-user-credits:{user_id}",
|
|
metadata_json={"legacy_credits": str(legacy), "migration_time": migration_time.isoformat()},
|
|
request_time=migration_time,
|
|
)
|
|
elif legacy < 0:
|
|
stats["negative_reset_users"] += 1
|
|
else:
|
|
stats["zero_users"] += 1
|
|
if do_commit:
|
|
await db.commit()
|
|
else:
|
|
await db.rollback()
|
|
except Exception:
|
|
await db.rollback()
|
|
raise
|
|
stats["migrated_credits"] = str(total.quantize(Decimal("0.01")))
|
|
stats["migration_time"] = migration_time.isoformat()
|
|
stats["mode"] = "commit" if do_commit else "dry-run"
|
|
print(json.dumps(stats, ensure_ascii=False, indent=2))
|
|
return 0
|
|
|
|
|
|
def main() -> int:
|
|
return asyncio.run(amain(sys.argv[1:]))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|