|
18 | 18 | logger = get_logger(__name__) |
19 | 19 |
|
20 | 20 |
|
| 21 | +class BillingError(Exception): |
| 22 | + """A charge could not be accepted; callers must roll back their transaction.""" |
| 23 | + |
| 24 | + def __init__(self, code: str, *, status: int = 503, details: Optional[dict] = None): |
| 25 | + super().__init__(code) |
| 26 | + self.code = code |
| 27 | + self.status = status |
| 28 | + self.details = details or {} |
| 29 | + |
| 30 | + |
21 | 31 | class BillingService: |
22 | 32 | """Billing and credit accounting service.""" |
23 | 33 |
|
@@ -595,6 +605,78 @@ def _grant_lifetime_monthly_credits_best_effort(self, cur, user_id: int): |
595 | 605 | # Best-effort; never break caller |
596 | 606 | pass |
597 | 607 |
|
| 608 | + def consume_in_transaction(self, cur, user_id: int, feature: str, reference_id: str) -> dict: |
| 609 | + """Deduct and record credits using the caller's transaction, without committing.""" |
| 610 | + enabled = self.is_billing_enabled() |
| 611 | + cost = max(0, int(self.get_feature_cost(feature) or 0)) |
| 612 | + cur.execute("SELECT credits FROM qd_users WHERE id = ? FOR UPDATE", (user_id,)) |
| 613 | + account = cur.fetchone() |
| 614 | + if not account: |
| 615 | + raise BillingError("BILLING_ACCOUNT_NOT_FOUND") |
| 616 | + balance = Decimal(str(account.get("credits") or 0)) |
| 617 | + charge = {"enabled": enabled, "feature": feature, "cost": cost, "charged": 0, |
| 618 | + "refunded": 0, "remaining": float(balance), "referenceId": reference_id, |
| 619 | + "transactionId": None, "refundTransactionId": None, "status": "free"} |
| 620 | + if reference_id: |
| 621 | + cur.execute("""SELECT id, amount, balance_after FROM qd_credits_log |
| 622 | + WHERE user_id = ? AND action = 'consume' AND feature = ? AND reference_id = ? |
| 623 | + ORDER BY id DESC LIMIT 1""", (user_id, feature, reference_id)) |
| 624 | + existing = cur.fetchone() |
| 625 | + if existing: |
| 626 | + amount = abs(int(existing["amount"])) |
| 627 | + return {**charge, "enabled": True, "cost": amount, "charged": amount, |
| 628 | + "transactionId": existing["id"], "status": "charged"} |
| 629 | + if not enabled or cost == 0: |
| 630 | + return charge |
| 631 | + cur.execute("""UPDATE qd_users SET credits = credits - ?, updated_at = NOW() |
| 632 | + WHERE id = ? AND credits >= ? RETURNING credits""", (cost, user_id, cost)) |
| 633 | + updated = cur.fetchone() |
| 634 | + if not updated: |
| 635 | + raise BillingError("INSUFFICIENT_CREDITS", status=402, details={ |
| 636 | + "feature": feature, "current": float(balance), "required": cost, |
| 637 | + "shortage": max(0, float(Decimal(cost) - balance)), |
| 638 | + }) |
| 639 | + remaining = float(updated["credits"]) |
| 640 | + cur.execute("""INSERT INTO qd_credits_log |
| 641 | + (user_id, action, amount, balance_after, feature, reference_id, remark, created_at) |
| 642 | + VALUES (?, 'consume', ?, ?, ?, ?, ?, ?) RETURNING id""", |
| 643 | + (user_id, -cost, remaining, feature, reference_id, |
| 644 | + f"Consume: {FEATURE_NAMES.get(feature, feature)}", datetime.now(timezone.utc))) |
| 645 | + transaction_id = cur.fetchone()["id"] |
| 646 | + return {**charge, "charged": cost, "remaining": remaining, |
| 647 | + "transactionId": transaction_id, "status": "charged"} |
| 648 | + |
| 649 | + def refund_in_transaction(self, cur, user_id: int, charge: dict) -> dict: |
| 650 | + """Refund the recorded debit once, atomically with the caller's job transition.""" |
| 651 | + if int(charge.get("charged") or 0) <= 0: |
| 652 | + return charge |
| 653 | + reference = str(charge["referenceId"]) |
| 654 | + feature = str(charge["feature"]) |
| 655 | + cur.execute("SELECT credits FROM qd_users WHERE id = ? FOR UPDATE", (user_id,)) |
| 656 | + if not cur.fetchone(): |
| 657 | + raise BillingError("BILLING_ACCOUNT_NOT_FOUND") |
| 658 | + cur.execute("""SELECT amount FROM qd_credits_log WHERE user_id = ? |
| 659 | + AND action = 'consume' AND feature = ? AND reference_id = ? |
| 660 | + ORDER BY id DESC LIMIT 1""", (user_id, feature, reference)) |
| 661 | + debit = cur.fetchone() |
| 662 | + if not debit: |
| 663 | + raise BillingError("BILLING_DEBIT_NOT_FOUND") |
| 664 | + amount = abs(int(debit["amount"])) |
| 665 | + cur.execute("""SELECT id, balance_after FROM qd_credits_log WHERE user_id = ? |
| 666 | + AND action = 'refund' AND reference_id = ? ORDER BY id DESC LIMIT 1""", (user_id, reference)) |
| 667 | + refund = cur.fetchone() |
| 668 | + if not refund: |
| 669 | + cur.execute("""UPDATE qd_users SET credits = credits + ?, updated_at = NOW() |
| 670 | + WHERE id = ? RETURNING credits""", (amount, user_id)) |
| 671 | + remaining = float(cur.fetchone()["credits"]) |
| 672 | + cur.execute("""INSERT INTO qd_credits_log |
| 673 | + (user_id, action, amount, balance_after, feature, reference_id, remark, created_at) |
| 674 | + VALUES (?, 'refund', ?, ?, ?, ?, ?, NOW()) RETURNING id""", |
| 675 | + (user_id, amount, remaining, feature, reference, "Automatic refund: agent job did not complete")) |
| 676 | + refund = {"id": cur.fetchone()["id"], "balance_after": remaining} |
| 677 | + return {**charge, "refunded": amount, "remaining": float(refund["balance_after"]), |
| 678 | + "refundTransactionId": refund["id"], "status": "refunded"} |
| 679 | + |
598 | 680 | def check_and_consume(self, user_id: int, feature: str, reference_id: str = '') -> Tuple[bool, str]: |
599 | 681 | """ |
600 | 682 | Check and consume credits for a feature. |
|
0 commit comments