fix:优化当前的项目

This commit is contained in:
Daniel
2026-04-28 18:36:38 +08:00
parent 04f26bdaaf
commit f47453a656
22 changed files with 3671 additions and 89 deletions

View File

@@ -2,6 +2,7 @@ from __future__ import annotations
import hashlib
import hmac
import json
import secrets
import sqlite3
import time
@@ -69,6 +70,12 @@ class UserStore:
pref_cols = self._table_columns(c, "user_prefs")
if "active_ai_model_id" not in pref_cols:
c.execute("ALTER TABLE user_prefs ADD COLUMN active_ai_model_id INTEGER")
if "subscriber_name" not in pref_cols:
c.execute("ALTER TABLE user_prefs ADD COLUMN subscriber_name TEXT NOT NULL DEFAULT ''")
if "subscriber_phone" not in pref_cols:
c.execute("ALTER TABLE user_prefs ADD COLUMN subscriber_phone TEXT NOT NULL DEFAULT ''")
if "shipping_address" not in pref_cols:
c.execute("ALTER TABLE user_prefs ADD COLUMN shipping_address TEXT NOT NULL DEFAULT ''")
c.execute(
"""
CREATE TABLE IF NOT EXISTS ai_models (
@@ -95,11 +102,62 @@ class UserStore:
vip_enabled INTEGER NOT NULL DEFAULT 0,
token_balance INTEGER NOT NULL DEFAULT 0,
total_consumed_tokens INTEGER NOT NULL DEFAULT 0,
seat_quota_credits INTEGER NOT NULL DEFAULT 1500,
seat_used_credits INTEGER NOT NULL DEFAULT 0,
seat_cycle TEXT NOT NULL DEFAULT '',
cycle_started_at INTEGER NOT NULL DEFAULT 0,
cycle_expires_at INTEGER NOT NULL DEFAULT 0,
updated_at INTEGER NOT NULL,
FOREIGN KEY(user_id) REFERENCES users(id)
)
"""
)
wallet_cols = self._table_columns(c, "user_wallets")
if "seat_quota_credits" not in wallet_cols:
c.execute("ALTER TABLE user_wallets ADD COLUMN seat_quota_credits INTEGER NOT NULL DEFAULT 25000")
if "seat_used_credits" not in wallet_cols:
c.execute("ALTER TABLE user_wallets ADD COLUMN seat_used_credits INTEGER NOT NULL DEFAULT 0")
if "seat_cycle" not in wallet_cols:
c.execute("ALTER TABLE user_wallets ADD COLUMN seat_cycle TEXT NOT NULL DEFAULT ''")
if "cycle_started_at" not in wallet_cols:
c.execute("ALTER TABLE user_wallets ADD COLUMN cycle_started_at INTEGER NOT NULL DEFAULT 0")
if "cycle_expires_at" not in wallet_cols:
c.execute("ALTER TABLE user_wallets ADD COLUMN cycle_expires_at INTEGER NOT NULL DEFAULT 0")
c.execute(
"""
CREATE TABLE IF NOT EXISTS recharge_orders (
id INTEGER PRIMARY KEY AUTOINCREMENT,
user_id INTEGER NOT NULL,
order_no TEXT NOT NULL UNIQUE,
channel TEXT NOT NULL DEFAULT '',
token_amount INTEGER NOT NULL DEFAULT 0,
amount_cny REAL NOT NULL DEFAULT 0.0,
status TEXT NOT NULL DEFAULT 'pending',
external_txn_id TEXT NOT NULL DEFAULT '',
meta_json TEXT NOT NULL DEFAULT '{}',
created_at INTEGER NOT NULL,
paid_at INTEGER,
FOREIGN KEY(user_id) REFERENCES users(id)
)
"""
)
c.execute(
"""
CREATE TABLE IF NOT EXISTS token_ledger (
id INTEGER PRIMARY KEY AUTOINCREMENT,
user_id INTEGER NOT NULL,
direction TEXT NOT NULL,
token_change INTEGER NOT NULL,
balance_after INTEGER NOT NULL,
kind TEXT NOT NULL DEFAULT '',
ref_type TEXT NOT NULL DEFAULT '',
ref_id TEXT NOT NULL DEFAULT '',
detail_json TEXT NOT NULL DEFAULT '{}',
created_at INTEGER NOT NULL,
FOREIGN KEY(user_id) REFERENCES users(id)
)
"""
)
ai_cols = self._table_columns(c, "ai_models")
if "image_model" not in ai_cols:
c.execute("ALTER TABLE ai_models ADD COLUMN image_model TEXT NOT NULL DEFAULT ''")
@@ -438,6 +496,57 @@ class UserStore:
return None
return {"id": int(row["id"]), "username": row["username"]}
def get_user_profile(self, user_id: int) -> dict:
with self._conn() as c:
row = c.execute(
"""
SELECT subscriber_name, subscriber_phone, shipping_address
FROM user_prefs
WHERE user_id=?
LIMIT 1
""",
(user_id,),
).fetchone()
if not row:
return {"subscriber_name": "", "subscriber_phone": "", "shipping_address": ""}
return {
"subscriber_name": (row["subscriber_name"] or "").strip(),
"subscriber_phone": (row["subscriber_phone"] or "").strip(),
"shipping_address": (row["shipping_address"] or "").strip(),
}
def save_user_profile(
self,
user_id: int,
*,
subscriber_name: str,
subscriber_phone: str,
shipping_address: str,
) -> dict:
now = int(time.time())
with self._conn() as c:
c.execute(
"""
INSERT INTO user_prefs(
user_id, active_wechat_account_id, active_ai_model_id,
subscriber_name, subscriber_phone, shipping_address, updated_at
) VALUES (?, NULL, NULL, ?, ?, ?, ?)
ON CONFLICT(user_id) DO UPDATE SET
subscriber_name=excluded.subscriber_name,
subscriber_phone=excluded.subscriber_phone,
shipping_address=excluded.shipping_address,
updated_at=excluded.updated_at
""",
(
user_id,
(subscriber_name or "").strip(),
(subscriber_phone or "").strip(),
(shipping_address or "").strip(),
now,
),
)
return self.get_user_profile(user_id)
def delete_user_logically(self, user_id: int, password: str, reset_code: str) -> bool:
now = int(time.time())
with self._conn() as c:
@@ -464,6 +573,8 @@ class UserStore:
c.execute("DELETE FROM user_prefs WHERE user_id=?", (user_id,))
c.execute("DELETE FROM wechat_bindings WHERE user_id=?", (user_id,))
c.execute("DELETE FROM user_wallets WHERE user_id=?", (user_id,))
c.execute("DELETE FROM recharge_orders WHERE user_id=?", (user_id,))
c.execute("DELETE FROM token_ledger WHERE user_id=?", (user_id,))
c.execute(
"UPDATE users SET deleted_at=?, username=username || '#deleted' || ? WHERE id=?",
(now, str(now), user_id),
@@ -472,14 +583,81 @@ class UserStore:
def _ensure_wallet_row(self, c: sqlite3.Connection, user_id: int) -> None:
now = int(time.time())
cycle = time.strftime("%Y-%m", time.localtime(now))
c.execute(
"""
INSERT OR IGNORE INTO user_wallets(user_id, vip_enabled, token_balance, total_consumed_tokens, updated_at)
VALUES (?, 0, 0, 0, ?)
INSERT OR IGNORE INTO user_wallets(
user_id, vip_enabled, token_balance, total_consumed_tokens,
seat_quota_credits, seat_used_credits, seat_cycle, cycle_started_at, cycle_expires_at, updated_at
)
VALUES (?, 0, 0, 0, 1500, 0, ?, 0, 0, ?)
""",
(user_id, now),
(user_id, cycle, now),
)
def _refresh_billing_cycle(self, c: sqlite3.Connection, user_id: int) -> None:
now = int(time.time())
row = c.execute(
"SELECT seat_cycle, cycle_expires_at, token_balance, seat_used_credits FROM user_wallets WHERE user_id=?",
(user_id,),
).fetchone()
current_cycle = (row["seat_cycle"] or "") if row else ""
expires_at = int(row["cycle_expires_at"] or 0) if row else 0
if expires_at > 0 and now >= expires_at:
c.execute(
"""
UPDATE user_wallets
SET seat_used_credits=0, token_balance=0, seat_cycle='', cycle_started_at=0, cycle_expires_at=0, updated_at=?
WHERE user_id=?
""",
(now, user_id),
)
return
# 兼容历史按自然月 seat_cycle 的老数据:若没有新周期字段,保留原行为
if expires_at <= 0:
paid = c.execute(
"""
SELECT paid_at
FROM recharge_orders
WHERE user_id=? AND status='paid' AND paid_at IS NOT NULL
ORDER BY paid_at DESC, id DESC
LIMIT 1
""",
(user_id,),
).fetchone()
if paid and int(paid["paid_at"] or 0) > 0:
start_at = int(paid["paid_at"])
new_expires = start_at + 30 * 24 * 3600
if now >= new_expires:
c.execute(
"""
UPDATE user_wallets
SET seat_used_credits=0, token_balance=0, seat_cycle='', cycle_started_at=0, cycle_expires_at=0, updated_at=?
WHERE user_id=?
""",
(now, user_id),
)
else:
c.execute(
"""
UPDATE user_wallets
SET seat_cycle=?, cycle_started_at=?, cycle_expires_at=?, updated_at=?
WHERE user_id=?
""",
(time.strftime("%Y-%m", time.localtime(start_at)), start_at, new_expires, now, user_id),
)
return
cycle = time.strftime("%Y-%m", time.localtime(now))
if current_cycle != cycle:
c.execute(
"""
UPDATE user_wallets
SET seat_used_credits=0, seat_cycle=?, updated_at=?
WHERE user_id=?
""",
(cycle, now, user_id),
)
def ensure_trial_tokens(self, user_id: int, trial_tokens: int) -> dict:
amount = max(0, int(trial_tokens))
now = int(time.time())
@@ -491,31 +669,63 @@ class UserStore:
).fetchone()
current = int(row["token_balance"] or 0) if row else 0
if current <= 0 and amount > 0:
new_balance = amount
c.execute(
"""
UPDATE user_wallets
SET vip_enabled=1, token_balance=?, updated_at=?
WHERE user_id=?
""",
(amount, now, user_id),
(new_balance, now, user_id),
)
c.execute(
"""
INSERT INTO token_ledger(
user_id, direction, token_change, balance_after, kind, ref_type, ref_id, detail_json, created_at
) VALUES (?, 'in', ?, ?, 'trial_grant', 'system', '', '{}', ?)
""",
(user_id, amount, new_balance, now),
)
return self.get_vip_status(user_id)
def get_vip_status(self, user_id: int) -> dict:
with self._conn() as c:
self._ensure_wallet_row(c, user_id)
self._refresh_billing_cycle(c, user_id)
row = c.execute(
"""
SELECT vip_enabled, token_balance, total_consumed_tokens, updated_at
SELECT
vip_enabled, token_balance, total_consumed_tokens,
seat_quota_credits, seat_used_credits, seat_cycle, cycle_started_at, cycle_expires_at, updated_at
FROM user_wallets
WHERE user_id=?
""",
(user_id,),
).fetchone()
seat_quota = int(row["seat_quota_credits"] or 0) if row else 0
seat_used = int(row["seat_used_credits"] or 0) if row else 0
seat_remaining = max(0, seat_quota - seat_used)
shared_credits = int(row["token_balance"] or 0) if row else 0
cycle_started_at = int(row["cycle_started_at"] or 0) if row else 0
cycle_expires_at = int(row["cycle_expires_at"] or 0) if row else 0
now = int(time.time())
cycle_active = cycle_expires_at > now if cycle_expires_at > 0 else True
if not cycle_active:
seat_remaining = 0
shared_credits = 0
return {
"vip_enabled": bool(int(row["vip_enabled"] or 0)) if row else False,
"token_balance": int(row["token_balance"] or 0) if row else 0,
"token_balance": shared_credits,
"total_consumed_tokens": int(row["total_consumed_tokens"] or 0) if row else 0,
"seat_quota_credits": seat_quota,
"seat_used_credits": seat_used,
"seat_remaining_credits": seat_remaining,
"shared_credits": shared_credits,
"total_available_credits": seat_remaining + shared_credits,
"seat_cycle": (row["seat_cycle"] or "") if row else "",
"cycle_started_at": cycle_started_at,
"cycle_expires_at": cycle_expires_at,
"cycle_active": cycle_active,
"updated_at": int(row["updated_at"] or 0) if row else 0,
}
@@ -529,45 +739,328 @@ class UserStore:
)
return self.get_vip_status(user_id)
def recharge_tokens(self, user_id: int, tokens: int) -> dict:
def recharge_tokens(
self,
user_id: int,
tokens: int,
*,
kind: str = "manual_recharge",
ref_type: str = "",
ref_id: str = "",
detail: dict | None = None,
cycle_start_at: int | None = None,
cycle_days: int = 30,
) -> dict:
add = max(0, int(tokens))
now = int(time.time())
with self._conn() as c:
self._ensure_wallet_row(c, user_id)
self._refresh_billing_cycle(c, user_id)
row = c.execute("SELECT token_balance FROM user_wallets WHERE user_id=?", (user_id,)).fetchone()
prev = int(row["token_balance"] or 0) if row else 0
new_balance = prev + add
start_at = int(cycle_start_at or 0)
if start_at > 0:
expires_at = start_at + max(1, int(cycle_days)) * 24 * 3600
c.execute(
"""
UPDATE user_wallets
SET token_balance=token_balance + ?, vip_enabled=1, seat_used_credits=0, seat_cycle=?, cycle_started_at=?, cycle_expires_at=?, updated_at=?
WHERE user_id=?
""",
(add, time.strftime("%Y-%m", time.localtime(start_at)), start_at, expires_at, now, user_id),
)
else:
c.execute(
"""
UPDATE user_wallets
SET token_balance=token_balance + ?, vip_enabled=1, updated_at=?
WHERE user_id=?
""",
(add, now, user_id),
)
c.execute(
"""
UPDATE user_wallets
SET token_balance=token_balance + ?, vip_enabled=1, updated_at=?
WHERE user_id=?
INSERT INTO token_ledger(
user_id, direction, token_change, balance_after, kind, ref_type, ref_id, detail_json, created_at
) VALUES (?, 'in', ?, ?, ?, ?, ?, ?, ?)
""",
(add, now, user_id),
(
user_id,
add,
new_balance,
kind,
ref_type,
ref_id or "",
json.dumps(detail or {}, ensure_ascii=True),
now,
),
)
return self.get_vip_status(user_id)
def consume_tokens(self, user_id: int, tokens: int) -> tuple[bool, int]:
def consume_tokens(
self,
user_id: int,
tokens: int,
*,
kind: str = "usage",
ref_type: str = "",
ref_id: str = "",
detail: dict | None = None,
) -> tuple[bool, int]:
cost = max(0, int(tokens))
now = int(time.time())
with self._conn() as c:
self._ensure_wallet_row(c, user_id)
self._refresh_billing_cycle(c, user_id)
row = c.execute(
"SELECT token_balance FROM user_wallets WHERE user_id=?",
"SELECT token_balance, seat_quota_credits, seat_used_credits FROM user_wallets WHERE user_id=?",
(user_id,),
).fetchone()
balance = int(row["token_balance"] or 0) if row else 0
shared_balance = int(row["token_balance"] or 0) if row else 0
seat_quota = int(row["seat_quota_credits"] or 0) if row else 0
seat_used = int(row["seat_used_credits"] or 0) if row else 0
seat_remaining = max(0, seat_quota - seat_used)
if cost <= 0:
return True, balance
if balance < cost:
return False, balance
new_balance = balance - cost
return True, seat_remaining + shared_balance
use_from_seat = min(seat_remaining, cost)
need_shared = cost - use_from_seat
if shared_balance < need_shared:
return False, seat_remaining + shared_balance
new_shared = shared_balance - need_shared
new_seat_used = seat_used + use_from_seat
c.execute(
"""
UPDATE user_wallets
SET token_balance=?, total_consumed_tokens=total_consumed_tokens + ?, updated_at=?
SET token_balance=?, seat_used_credits=?, total_consumed_tokens=total_consumed_tokens + ?, updated_at=?
WHERE user_id=?
""",
(new_balance, cost, now, user_id),
(new_shared, new_seat_used, cost, now, user_id),
)
return True, new_balance
c.execute(
"""
INSERT INTO token_ledger(
user_id, direction, token_change, balance_after, kind, ref_type, ref_id, detail_json, created_at
) VALUES (?, 'out', ?, ?, ?, ?, ?, ?, ?)
""",
(
user_id,
cost,
max(0, seat_quota - new_seat_used) + new_shared,
kind,
ref_type,
ref_id or "",
json.dumps(
{
**(detail or {}),
"credit_source": {"seat": use_from_seat, "shared": need_shared},
},
ensure_ascii=True,
),
now,
),
)
return True, max(0, seat_quota - new_seat_used) + new_shared
def create_recharge_order(
self,
user_id: int,
order_no: str,
channel: str,
token_amount: int,
amount_cny: float,
meta: dict | None = None,
) -> dict:
now = int(time.time())
with self._conn() as c:
c.execute(
"""
INSERT INTO recharge_orders(
user_id, order_no, channel, token_amount, amount_cny, status, external_txn_id, meta_json, created_at, paid_at
) VALUES (?, ?, ?, ?, ?, 'pending', '', ?, ?, NULL)
""",
(
user_id,
order_no,
channel or "",
int(token_amount),
float(amount_cny),
json.dumps(meta or {}, ensure_ascii=True),
now,
),
)
return {
"order_no": order_no,
"channel": channel,
"token_amount": int(token_amount),
"amount_cny": float(amount_cny),
"status": "pending",
"created_at": now,
}
def mark_recharge_order_paid(
self,
user_id: int,
order_no: str,
paid_amount_cny: float,
external_txn_id: str = "",
meta: dict | None = None,
) -> tuple[bool, str]:
now = int(time.time())
with self._conn() as c:
row = c.execute(
"""
SELECT user_id, token_amount, amount_cny, status
FROM recharge_orders
WHERE order_no=?
""",
(order_no,),
).fetchone()
if not row:
return False, "订单不存在"
if int(row["user_id"]) != int(user_id):
return False, "订单无权限"
if (row["status"] or "") == "paid":
return True, "already_paid"
if float(paid_amount_cny or 0.0) + 1e-9 < float(row["amount_cny"] or 0.0):
return False, "支付金额不足"
c.execute(
"""
UPDATE recharge_orders
SET status='paid', external_txn_id=?, paid_at=?, meta_json=?
WHERE order_no=?
""",
(
external_txn_id or "",
now,
json.dumps(meta or {}, ensure_ascii=True),
order_no,
),
)
self.recharge_tokens(
user_id,
int(row["token_amount"] or 0),
kind="paid_recharge",
ref_type="order",
ref_id=order_no,
detail={"paid_amount_cny": float(paid_amount_cny or 0.0), "external_txn_id": external_txn_id or ""},
cycle_start_at=now,
cycle_days=30,
)
return True, "ok"
def list_recharge_orders(self, user_id: int, limit: int = 50) -> list[dict]:
with self._conn() as c:
now = int(time.time())
expire_before = now - 15 * 60
c.execute(
"""
UPDATE recharge_orders
SET status='cancelled'
WHERE user_id=? AND status='pending' AND created_at<=?
""",
(user_id, expire_before),
)
rows = c.execute(
"""
SELECT order_no, channel, token_amount, amount_cny, status, external_txn_id, created_at, paid_at, meta_json
FROM recharge_orders
WHERE user_id=?
ORDER BY id DESC
LIMIT ?
""",
(user_id, max(1, min(int(limit), 200))),
).fetchall()
return [
{
"order_no": r["order_no"] or "",
"channel": r["channel"] or "",
"token_amount": int(r["token_amount"] or 0),
"amount_cny": float(r["amount_cny"] or 0.0),
"status": r["status"] or "",
"external_txn_id": r["external_txn_id"] or "",
"created_at": int(r["created_at"] or 0),
"paid_at": int(r["paid_at"] or 0) if r["paid_at"] else None,
"meta": json.loads(r["meta_json"] or "{}"),
}
for r in rows
]
def get_recharge_order(self, user_id: int, order_no: str) -> dict | None:
now = int(time.time())
with self._conn() as c:
expire_before = now - 15 * 60
c.execute(
"""
UPDATE recharge_orders
SET status='cancelled'
WHERE user_id=? AND status='pending' AND created_at<=?
""",
(user_id, expire_before),
)
row = c.execute(
"""
SELECT order_no, channel, token_amount, amount_cny, status, external_txn_id, created_at, paid_at, meta_json
FROM recharge_orders
WHERE user_id=? AND order_no=?
LIMIT 1
""",
(user_id, order_no),
).fetchone()
if not row:
return None
try:
meta = json.loads(row["meta_json"] or "{}")
except Exception:
meta = {}
return {
"order_no": row["order_no"] or "",
"channel": row["channel"] or "",
"token_amount": int(row["token_amount"] or 0),
"amount_cny": float(row["amount_cny"] or 0.0),
"status": row["status"] or "",
"external_txn_id": row["external_txn_id"] or "",
"created_at": int(row["created_at"] or 0),
"paid_at": int(row["paid_at"] or 0) if row["paid_at"] else None,
"meta": meta,
}
def get_recharge_order_user_id(self, order_no: str) -> int | None:
with self._conn() as c:
row = c.execute("SELECT user_id FROM recharge_orders WHERE order_no=?", (order_no,)).fetchone()
return int(row["user_id"]) if row and row["user_id"] else None
def list_token_ledger(self, user_id: int, limit: int = 100) -> list[dict]:
with self._conn() as c:
rows = c.execute(
"""
SELECT direction, token_change, balance_after, kind, ref_type, ref_id, detail_json, created_at
FROM token_ledger
WHERE user_id=?
ORDER BY id DESC
LIMIT ?
""",
(user_id, max(1, min(int(limit), 500))),
).fetchall()
out: list[dict] = []
for r in rows:
try:
detail = json.loads(r["detail_json"] or "{}")
except Exception:
detail = {}
out.append(
{
"direction": r["direction"] or "",
"token_change": int(r["token_change"] or 0),
"balance_after": int(r["balance_after"] or 0),
"kind": r["kind"] or "",
"ref_type": r["ref_type"] or "",
"ref_id": r["ref_id"] or "",
"detail": detail,
"created_at": int(r["created_at"] or 0),
}
)
return out
def save_wechat_binding(
self,
@@ -887,6 +1380,41 @@ class UserStore:
)
return True
def update_active_ai_image_model(self, user_id: int, image_model: str) -> bool:
now = int(time.time())
name = (image_model or "").strip()
with self._conn() as c:
pref = c.execute(
"SELECT active_ai_model_id FROM user_prefs WHERE user_id=?",
(user_id,),
).fetchone()
aid = int(pref["active_ai_model_id"]) if pref and pref["active_ai_model_id"] else None
if not aid:
row = c.execute(
"SELECT id FROM ai_models WHERE user_id=? ORDER BY updated_at DESC, id DESC LIMIT 1",
(user_id,),
).fetchone()
aid = int(row["id"]) if row else None
if not aid:
return False
c.execute(
"UPDATE ai_models SET image_model=?, updated_at=? WHERE id=? AND user_id=?",
(name, now, aid, user_id),
)
if c.total_changes <= 0:
return False
c.execute(
"""
INSERT INTO user_prefs(user_id, active_ai_model_id, updated_at)
VALUES (?, ?, ?)
ON CONFLICT(user_id) DO UPDATE SET
active_ai_model_id=excluded.active_ai_model_id,
updated_at=excluded.updated_at
""",
(user_id, aid, now),
)
return True
def get_active_ai_model(self, user_id: int) -> dict | None:
with self._conn() as c:
pref = c.execute(