# -*- coding: utf-8 -*- """预扣账本（开账 / 改名 / 补扣 / 结算 / 对账）。 （自 auth_system.py 拆出，行为与原实现一致。） """ import time from typing import Any from .db import ( _DB_LOCK, _q, _x, get_conn, ) from .logs import ( add_log, ) # ============ 计费账本（实扣、交付、退款、终态的持久化唯一记录） ============ def ledger_open(task_id: str, user_id: int, reserved: int) -> None: """任务计费开账：记录本任务实际预扣次数（admin/-1 用户为 0）。 task_id 唯一；重复 task_id 由任务系统在提交侧拒绝，这里幂等保护。""" with _DB_LOCK: conn = get_conn() try: conn.execute( 'INSERT OR IGNORE INTO billing_ledger (task_id, user_id, reserved, status, created_at) ' 'VALUES (?, ?, ?, \'open\', ?)', (str(task_id), int(user_id), max(0, int(reserved)), int(time.time()))) conn.commit() finally: conn.close() def reserve_and_open(uid: int, count: int, ledger_key: str) -> int | None: """预扣与开账同一事务（原两步之间存在崩溃丢账窗口）。 足额：扣减+开账并返回剩余；不限/管理员：开账(reserved=0)返回 -1；不足：返回 None（不开账）。""" count = max(1, int(count)) key = str(ledger_key) with _DB_LOCK: conn = get_conn() try: row = conn.execute('SELECT role, duration_seconds FROM users WHERE id = ?', (uid,)).fetchone() if not row: return None now = int(time.time()) if row['role'] == 'admin': conn.execute('UPDATE users SET last_active = ? WHERE id = ?', (now, uid)) conn.execute( 'INSERT OR IGNORE INTO billing_ledger (task_id, user_id, reserved, status, created_at) ' 'VALUES (?, ?, 0, \'open\', ?)', (key, int(uid), now)) conn.commit() return -1 cur = int(row['duration_seconds'] or 0) if cur == -1: conn.execute('UPDATE users SET last_active = ? WHERE id = ?', (now, uid)) conn.execute( 'INSERT OR IGNORE INTO billing_ledger (task_id, user_id, reserved, status, created_at) ' 'VALUES (?, ?, 0, \'open\', ?)', (key, int(uid), now)) conn.commit() return -1 if cur < count: return None conn.execute('UPDATE users SET duration_seconds = ?, last_active = ? WHERE id = ?', (cur - count, now, uid)) conn.execute( 'INSERT OR IGNORE INTO billing_ledger (task_id, user_id, reserved, status, created_at) ' 'VALUES (?, ?, ?, \'open\', ?)', (key, int(uid), count, now)) conn.commit() return cur - count finally: conn.close() def ledger_rename(old_key: str, new_key: str) -> bool: """把请求级临时账本行改名为正式 task_id（任务接手账目）。 old 不存在/已关闭/new 已存在 → False。""" with _DB_LOCK: conn = get_conn() try: dup = conn.execute('SELECT 1 FROM billing_ledger WHERE task_id = ?', (str(new_key),)).fetchone() if dup: return False cur = conn.execute( 'UPDATE billing_ledger SET task_id = ? WHERE task_id = ? AND status = \'open\'', (str(new_key), str(old_key))) conn.commit() return cur.rowcount == 1 finally: conn.close() def ledger_supplement(task_id: str, extra: int) -> bool: """整台播客等"先扣 1 后按实际数量补扣"场景： 原子补扣用户余额并同步累加账本 reserved。任一步失败返回 False（调用方必须让任务失败， 不允许无账本继续交付）。余额不足同样返回 False。""" task_id = str(task_id) extra = int(extra) if extra <= 0: return True with _DB_LOCK: conn = get_conn() try: row = conn.execute('SELECT user_id, status FROM billing_ledger WHERE task_id = ?', (task_id,)).fetchone() if not row or row['status'] != 'open': return False urow = conn.execute('SELECT role, duration_seconds FROM users WHERE id = ?', (row['user_id'],)).fetchone() if not urow or urow['role'] == 'admin' or int(urow['duration_seconds'] or 0) == -1: return True # 不限用户：无需补扣，账本也不加 cur_credits = int(urow['duration_seconds'] or 0) if cur_credits < extra: return False # 余额不足：任务必须失败（不允许扣到 0 硬做） conn.execute('UPDATE users SET duration_seconds = duration_seconds - ? WHERE id = ?', (extra, row['user_id'])) conn.execute('UPDATE billing_ledger SET reserved = reserved + ? WHERE task_id = ?', (extra, task_id)) conn.commit() return True except Exception: conn.rollback() return False finally: conn.close() def ledger_get(task_id: str) -> dict[str, Any] | None: rows = _q('SELECT * FROM billing_ledger WHERE task_id = ?', (str(task_id),)) return dict(rows[0]) if rows else None def ledger_settle(task_id: str, delivered: int, final_status: str) -> int: """任务计费结算（幂等）：按账本实扣与实际交付计算应退次数，原子退款并落终态。 - delivered：本次任务实际交付给用户的数量（单曲/单集 0 或 1；批量为成功入包数） - 重复调用/重启后再次调用：账本已非 open 直接返回 0，不重复退款 - admin/-1 提交（reserved=0）：退款自然为 0（不看结算时的账户状态） - 以账本 reserved 为准计算应退；口径更正：-1（不限次）账户不落余额 （加次数会破坏 -1 哨兵值），账本 refunded 记录应付口径，作废可接受 返回实际退款次数（落余额的部分；-1/admin 场景为 0）。""" task_id = str(task_id) delivered = max(0, int(delivered)) with _DB_LOCK: conn = get_conn() try: row = conn.execute('SELECT * FROM billing_ledger WHERE task_id = ?', (task_id,)).fetchone() if not row or row['status'] != 'open': return 0 refund = max(0, int(row['reserved']) - delivered) # 退款与账本终态同事务落库：要么都成，要么都不成。 # 注意 WHERE 必须排除 admin/-1：向 -1（不限次）账户加次数会把 -1 哨兵值破坏成 # -1+N，向 admin 加无意义——因此任务运行期间账户升级为 -1 时退款不落余额， # 账本 refunded 记录的是"按账本应付"口径（该用户已有无限次数，作废可接受）； # 原注释"照样退"与实现相反，已更正为此口径。 if refund > 0: conn.execute('UPDATE users SET duration_seconds = duration_seconds + ? ' 'WHERE id = ? AND role != \'admin\' AND duration_seconds != -1', (refund, row['user_id'])) conn.execute('UPDATE billing_ledger SET delivered = ?, refunded = ?, status = ?, ' 'settled_at = ? WHERE task_id = ?', (delivered, refund, str(final_status)[:20], int(time.time()), task_id)) conn.commit() return refund finally: conn.close() def ledger_recover_orphans(max_age_seconds: int = 2 * 3600) -> int: """启动对账：任务表在内存，进程重启后账本中仍 open 的行代表 "已扣款但任务已丢失"——退回全部实扣。仅处理超过宽限期的行， 避免误伤同库其他存活进程的进行中任务。返回退款笔数。""" now = int(time.time()) recovered = 0 rows = _q('SELECT task_id FROM billing_ledger WHERE status = \'open\' AND created_at <= ?', (now - int(max_age_seconds),)) for r in rows: try: refunded = ledger_settle(r['task_id'], 0, 'lost_restart') if refunded: add_log('success', 'system', 'system', '重启对账退款', f'任务 {r["task_id"]} 因进程重启丢失，退回 {refunded} 次') else: ledger_settle(r['task_id'], 0, 'lost_restart') recovered += 1 except Exception: pass return recovered def ledger_cleanup_settled(max_age_seconds: int = 90 * 86400) -> int: """清理已结算的历史账本行（每任务一行、只增不删会让表无限膨胀； settled 状态行无对账价值，保留 90 天足够追溯）""" try: return _x('DELETE FROM billing_ledger WHERE status != \'open\' AND settled_at <= ?', (int(time.time()) - int(max_age_seconds),)) except Exception: return 0