from __future__ import annotations import hashlib import json from pathlib import Path from threading import RLock from typing import Any, Callable, TypeVar from sqlalchemy import Column, String, Text, create_engine, Integer, event, inspect, text from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker from services.storage.base import StorageBackend Base = declarative_base() ResultT = TypeVar("ResultT") class AccountModel(Base): """账号数据模型""" __tablename__ = "accounts" id = Column(Integer, primary_key=True, autoincrement=True) access_token = Column(String(2048), unique=True, nullable=False, index=True) owner_id = Column(String(255), nullable=True, index=True) shared = Column(Integer, nullable=False, default=0, index=True) status = Column(String(64), nullable=True, index=True) data = Column(Text, nullable=False) # JSON 格式存储完整账号数据 class AuthKeyModel(Base): """鉴权密钥数据模型""" __tablename__ = "auth_keys" id = Column(Integer, primary_key=True, autoincrement=True) key_id = Column(String(255), unique=True, nullable=False, index=True) data = Column(Text, nullable=False) class ShopStateModel(Base): __tablename__ = "shop_state" key = Column(String(64), primary_key=True) data = Column(Text, nullable=False) class RedeemCodeModel(Base): __tablename__ = "redeem_codes" id = Column(Integer, primary_key=True, autoincrement=True) code_id = Column(String(64), unique=True, nullable=False, index=True) code_hash = Column(String(128), unique=True, nullable=False, index=True) batch_id = Column(String(64), nullable=True, index=True) redeemed_by = Column(String(255), nullable=True, index=True) status = Column(String(32), nullable=False, index=True) data = Column(Text, nullable=False) class ShopLedgerModel(Base): __tablename__ = "shop_ledger" id = Column(Integer, primary_key=True, autoincrement=True) ledger_id = Column(String(64), unique=True, nullable=False, index=True) created_at = Column(String(64), nullable=True, index=True) user_id = Column(String(255), nullable=True, index=True) data = Column(Text, nullable=False) class DatabaseStorageBackend(StorageBackend): """数据库存储后端（支持 SQLite、PostgreSQL、MySQL 等）""" def __init__(self, database_url: str, *, table_scope: str = "all"): self.database_url = database_url self.table_scope = table_scope if table_scope in {"all", "accounts", "users"} else "all" connect_args: dict[str, Any] = {} if database_url.startswith("sqlite:///"): db_path = database_url.removeprefix("sqlite:///") if db_path and db_path != ":memory:": Path(db_path).parent.mkdir(parents=True, exist_ok=True) connect_args["check_same_thread"] = False connect_args["timeout"] = 30 self.engine = create_engine( database_url, pool_pre_ping=True, # 自动检测连接是否有效 pool_recycle=3600, # 1小时回收连接 connect_args=connect_args, ) if database_url.startswith("sqlite:///"): @event.listens_for(self.engine, "connect") def _configure_sqlite_connection(dbapi_connection: Any, _connection_record: Any) -> None: cursor = dbapi_connection.cursor() try: cursor.execute("PRAGMA busy_timeout=30000") cursor.execute("PRAGMA synchronous=FULL") cursor.execute("PRAGMA journal_mode=DELETE") finally: cursor.close() tables = None if self.table_scope == "accounts": tables = [AccountModel.__table__] elif self.table_scope == "users": tables = [ AuthKeyModel.__table__, ShopStateModel.__table__, RedeemCodeModel.__table__, ShopLedgerModel.__table__, ] Base.metadata.create_all(self.engine, tables=tables) self.Session = sessionmaker(bind=self.engine) self._lock = RLock() if self.table_scope in {"all", "accounts"}: self._migrate_account_query_columns() if self.table_scope in {"all", "users"}: self._migrate_legacy_redeem_codes() self._migrate_legacy_shop_ledger() def load_accounts(self) -> list[dict[str, Any]]: """从数据库加载账号数据""" session = self.Session() try: accounts = [] for row in session.query(AccountModel).all(): try: account_data = json.loads(row.data) if isinstance(account_data, dict): accounts.append(account_data) except json.JSONDecodeError: continue return accounts finally: session.close() @staticmethod def _account_owner_id(item: dict[str, Any]) -> str: return str(item.get("owner_id") or "").strip() @staticmethod def _account_shared(item: dict[str, Any]) -> int: return 1 if bool(item.get("shared")) else 0 @staticmethod def _account_status(item: dict[str, Any]) -> str: return str(item.get("status") or "").strip() def _account_from_row(self, row: AccountModel) -> dict[str, Any] | None: try: account_data = json.loads(row.data) except json.JSONDecodeError: return None return account_data if isinstance(account_data, dict) else None def _apply_account_row(self, row: AccountModel, item: dict[str, Any]) -> None: row.access_token = str(item.get("access_token") or "").strip() row.owner_id = self._account_owner_id(item) or None row.shared = self._account_shared(item) row.status = self._account_status(item) or None row.data = json.dumps(item, ensure_ascii=False) def _load_account_rows(self, query: Any) -> list[dict[str, Any]]: accounts: list[dict[str, Any]] = [] for row in query.all(): account_data = self._account_from_row(row) if account_data is not None: accounts.append(account_data) return accounts def load_public_accounts(self) -> list[dict[str, Any]]: session = self.Session() try: return self._load_account_rows( session.query(AccountModel).filter( (AccountModel.owner_id.is_(None)) | (AccountModel.owner_id == "") ) ) finally: session.close() def load_shared_accounts(self) -> list[dict[str, Any]]: session = self.Session() try: return self._load_account_rows( session.query(AccountModel).filter( AccountModel.owner_id.isnot(None), AccountModel.owner_id != "", AccountModel.shared == 1, ) ) finally: session.close() def load_accounts_by_owner(self, owner_id: str) -> list[dict[str, Any]]: normalized_owner = str(owner_id or "").strip() if not normalized_owner: return [] session = self.Session() try: return self._load_account_rows(session.query(AccountModel).filter_by(owner_id=normalized_owner)) finally: session.close() def load_account(self, access_token: str) -> dict[str, Any] | None: normalized_token = str(access_token or "").strip() if not normalized_token: return None session = self.Session() try: row = session.query(AccountModel).filter_by(access_token=normalized_token).first() return self._account_from_row(row) if row is not None else None finally: session.close() def upsert_accounts(self, accounts: list[dict[str, Any]]) -> None: with self._lock: session = self.Session() try: self._acquire_transaction_lock(session, "accounts") for item in accounts: if not isinstance(item, dict): continue access_token = str(item.get("access_token") or "").strip() if not access_token: continue row = session.query(AccountModel).filter_by(access_token=access_token).first() if row is None: row = AccountModel(access_token=access_token, data="{}") session.add(row) self._apply_account_row(row, item) session.commit() except Exception as e: session.rollback() raise e finally: session.close() def delete_accounts_by_tokens( self, tokens: list[str], *, owner_id: str = "", admin: bool = True, ) -> None: normalized_tokens = [str(token or "").strip() for token in dict.fromkeys(tokens) if str(token or "").strip()] if not normalized_tokens: return normalized_owner = str(owner_id or "").strip() with self._lock: session = self.Session() try: self._acquire_transaction_lock(session, "accounts") query = session.query(AccountModel).filter(AccountModel.access_token.in_(normalized_tokens)) if not admin: query = query.filter_by(owner_id=normalized_owner) query.delete(synchronize_session=False) session.commit() except Exception as e: session.rollback() raise e finally: session.close() def save_accounts(self, accounts: list[dict[str, Any]]) -> None: """保存账号数据到数据库""" with self._lock: session = self.Session() try: self._acquire_transaction_lock(session, "accounts") session.query(AccountModel).delete() for item in accounts: if not isinstance(item, dict): continue access_token = str(item.get("access_token") or "").strip() if not access_token: continue row = AccountModel(access_token=access_token, data="{}") self._apply_account_row(row, item) session.add(row) session.commit() except Exception as e: session.rollback() raise e finally: session.close() def load_auth_keys(self) -> list[dict[str, Any]]: """从数据库加载鉴权密钥数据""" return self._load_rows(AuthKeyModel) def save_auth_keys(self, auth_keys: list[dict[str, Any]]) -> None: """保存鉴权密钥数据到数据库""" with self._lock: self._save_rows(AuthKeyModel, auth_keys, "id", "key_id") def merge_auth_keys(self, changed_items: list[dict[str, Any]], deleted_ids: list[str]) -> None: """Apply auth-key changes without rewriting unchanged rows. This is safer for shared PostgreSQL deployments where multiple Spaces can update different users at the same time. """ with self._lock: session = self.Session() try: self._acquire_transaction_lock(session, "auth_keys") for key_id in {str(item or "").strip() for item in deleted_ids if str(item or "").strip()}: session.query(AuthKeyModel).filter_by(key_id=key_id).delete() for item in changed_items: if not isinstance(item, dict): continue key_id = str(item.get("id") or "").strip() if not key_id: continue payload = json.dumps(item, ensure_ascii=False) row = session.query(AuthKeyModel).filter_by(key_id=key_id).first() if row is None: session.add(AuthKeyModel(key_id=key_id, data=payload)) else: row.data = payload session.commit() except Exception as e: session.rollback() raise e finally: session.close() def update_auth_keys_locked( self, mutator: Callable[[list[dict[str, Any]]], tuple[list[dict[str, Any]], ResultT]], ) -> ResultT: """Load, mutate, and save auth keys under one database transaction lock.""" with self._lock: session = self.Session() try: self._acquire_transaction_lock(session, "auth_keys") rows = session.query(AuthKeyModel).all() items: list[dict[str, Any]] = [] for row in rows: try: item_data = json.loads(row.data) except json.JSONDecodeError: continue if isinstance(item_data, dict): items.append(item_data) next_items, result = mutator(items) normalized_items = [item for item in next_items if isinstance(item, dict)] next_ids = {str(item.get("id") or "").strip() for item in normalized_items} next_ids.discard("") for row in rows: if str(row.key_id or "").strip() not in next_ids: session.delete(row) existing_by_id = {str(row.key_id or "").strip(): row for row in rows} for item in normalized_items: key_id = str(item.get("id") or "").strip() if not key_id: continue payload = json.dumps(item, ensure_ascii=False) row = existing_by_id.get(key_id) if row is None: session.add(AuthKeyModel(key_id=key_id, data=payload)) elif row.data != payload: # ⚠️⚠️ 只写真正变化的行。之前是无条件赋值，等于每次「改一条」都 # 把整张表（几百上千行）标记为 dirty 后逐行 UPDATE + commit。 # 图片任务线程并发结算积分时，这一步会把 SQLite 写锁占住几百毫秒， # 而事件循环线程也在等这把锁 ⇒ 整站 HTTP 卡死（2026-10-19 事故）。 row.data = payload session.commit() return result except Exception as e: session.rollback() raise e finally: session.close() @staticmethod def _loads_dict(value: str | None) -> dict[str, Any]: try: data = json.loads(value or "{}") except json.JSONDecodeError: return {} return data if isinstance(data, dict) else {} @staticmethod def _compact_shop_state(state: dict[str, Any]) -> dict[str, Any]: compact = dict(state if isinstance(state, dict) else {}) codes = compact.get("codes") if isinstance(compact.get("codes"), list) else [] compact["codes"] = [ item for item in codes if isinstance(item, dict) and not (str(item.get("id") or "").strip() and str(item.get("code_hash") or "").strip()) ] # Redeem codes and the ledger are stored in their own indexed tables. # Keep only malformed legacy values here so an upgrade cannot discard # data it cannot yet identify. compact["ledger"] = [] return compact @staticmethod def _code_identity(item: dict[str, Any]) -> tuple[str, str]: return str(item.get("id") or "").strip(), str(item.get("code_hash") or "").strip() def _load_shop_state_from_session(self, session: Any, key: str = "default") -> dict[str, Any]: row = session.query(ShopStateModel).filter_by(key=str(key or "").strip() or "default").first() return self._loads_dict(row.data if row is not None else None) def _save_shop_state_to_session(self, session: Any, key: str, state: dict[str, Any]) -> None: normalized_key = str(key or "").strip() or "default" payload_state = self._compact_shop_state(state) if normalized_key == "default" else state self._save_raw_shop_state_to_session(session, normalized_key, payload_state) def _save_raw_shop_state_to_session(self, session: Any, key: str, state: dict[str, Any]) -> None: normalized_key = str(key or "").strip() or "default" payload = json.dumps(state if isinstance(state, dict) else {}, ensure_ascii=False) row = session.query(ShopStateModel).filter_by(key=normalized_key).first() if row is None: session.add(ShopStateModel(key=normalized_key, data=payload)) else: row.data = payload def _redeem_code_from_row(self, row: RedeemCodeModel) -> dict[str, Any]: item = self._loads_dict(row.data) item.setdefault("id", row.code_id) item.setdefault("code_hash", row.code_hash) item.setdefault("status", row.status) if row.batch_id and not item.get("batch_id"): item["batch_id"] = row.batch_id if row.redeemed_by and not item.get("redeemed_by"): item["redeemed_by"] = row.redeemed_by return item def _load_redeem_codes_from_session(self, session: Any) -> list[dict[str, Any]]: items: list[dict[str, Any]] = [] for row in session.query(RedeemCodeModel).order_by(RedeemCodeModel.id.asc()).all(): item = self._redeem_code_from_row(row) code_id, code_hash = self._code_identity(item) if code_id and code_hash: items.append(item) return items def _apply_redeem_code_row(self, row: RedeemCodeModel, item: dict[str, Any]) -> None: code_id, code_hash = self._code_identity(item) batch_id = str(item.get("batch_id") or "").strip() or None redeemed_by = str(item.get("redeemed_by") or "").strip() or None status = str(item.get("status") or "active").strip() or "active" payload = json.dumps(item, ensure_ascii=False) if ( row.code_id == code_id and row.code_hash == code_hash and row.batch_id == batch_id and row.redeemed_by == redeemed_by and row.status == status and row.data == payload ): return row.code_id = code_id row.code_hash = code_hash row.batch_id = batch_id row.redeemed_by = redeemed_by row.status = status row.data = payload def _sync_redeem_codes_to_session(self, session: Any, codes: list[dict[str, Any]]) -> None: normalized_items: list[dict[str, Any]] = [] seen_hashes: set[str] = set() for item in codes: if not isinstance(item, dict): continue code_id, code_hash = self._code_identity(item) if not code_id or not code_hash or code_hash in seen_hashes: continue seen_hashes.add(code_hash) normalized_items.append(dict(item)) rows = session.query(RedeemCodeModel).all() existing_by_hash = {str(row.code_hash or "").strip(): row for row in rows} next_hashes = {str(item.get("code_hash") or "").strip() for item in normalized_items} for row in rows: if str(row.code_hash or "").strip() not in next_hashes: session.delete(row) for item in normalized_items: code_hash = str(item.get("code_hash") or "").strip() row = existing_by_hash.get(code_hash) if row is None: row = RedeemCodeModel( code_id=str(item.get("id") or "").strip(), code_hash=code_hash, status=str(item.get("status") or "active").strip() or "active", data="{}", ) session.add(row) self._apply_redeem_code_row(row, item) @staticmethod def _ledger_identity(item: dict[str, Any], position: int) -> str: ledger_id = str(item.get("id") or "").strip() if ledger_id: return ledger_id payload = json.dumps(item, ensure_ascii=False, sort_keys=True, separators=(",", ":")) digest = hashlib.sha256(f"{position}:{payload}".encode("utf-8")).hexdigest()[:24] return f"legacy-{digest}" def _normalize_ledger_items(self, ledger: list[dict[str, Any]]) -> list[dict[str, Any]]: normalized: list[dict[str, Any]] = [] seen_ids: set[str] = set() for position, item in enumerate(ledger): if not isinstance(item, dict): continue normalized_item = dict(item) ledger_id = self._ledger_identity(normalized_item, position) if ledger_id in seen_ids: duplicate_digest = hashlib.sha256( f"{ledger_id}:{position}".encode("utf-8") ).hexdigest()[:12] ledger_id = f"{ledger_id[:48]}-{duplicate_digest}" seen_ids.add(ledger_id) normalized_item["id"] = ledger_id normalized.append(normalized_item) return normalized def _ledger_from_row(self, row: ShopLedgerModel) -> dict[str, Any]: item = self._loads_dict(row.data) item.setdefault("id", row.ledger_id) if row.created_at and not item.get("time"): item["time"] = row.created_at if row.user_id and not item.get("user_id"): item["user_id"] = row.user_id return item def _load_shop_ledger_from_session(self, session: Any) -> list[dict[str, Any]]: return [ self._ledger_from_row(row) for row in session.query(ShopLedgerModel).order_by(ShopLedgerModel.id.asc()).all() ] def _apply_shop_ledger_row(self, row: ShopLedgerModel, item: dict[str, Any]) -> None: payload = json.dumps(item, ensure_ascii=False) ledger_id = str(item["id"]) created_at = str(item.get("time") or "").strip() or None user_id = str(item.get("user_id") or "").strip() or None if ( row.ledger_id == ledger_id and row.created_at == created_at and row.user_id == user_id and row.data == payload ): return row.ledger_id = ledger_id row.created_at = created_at row.user_id = user_id row.data = payload def _sync_shop_ledger_to_session(self, session: Any, ledger: list[dict[str, Any]]) -> None: normalized_items = self._normalize_ledger_items(ledger) rows = session.query(ShopLedgerModel).all() existing_by_id = {str(row.ledger_id or "").strip(): row for row in rows} next_ids = {str(item["id"]) for item in normalized_items} for row in rows: if str(row.ledger_id or "").strip() not in next_ids: session.delete(row) for item in normalized_items: ledger_id = str(item["id"]) row = existing_by_id.get(ledger_id) if row is None: row = ShopLedgerModel(ledger_id=ledger_id, data="{}") session.add(row) self._apply_shop_ledger_row(row, item) def _append_shop_ledger_to_session(self, session: Any, entries: list[dict[str, Any]]) -> None: """Append ledger entries without scanning or rewriting prior history.""" for item in self._normalize_ledger_items(entries): ledger_id = str(item["id"]) row = session.query(ShopLedgerModel).filter_by(ledger_id=ledger_id).first() if row is None: row = ShopLedgerModel(ledger_id=ledger_id, data="{}") session.add(row) self._apply_shop_ledger_row(row, item) def _load_default_shop_state_from_session(self, session: Any) -> dict[str, Any]: state = self._load_shop_state_from_session(session, "default") legacy_codes = state.get("codes") if isinstance(state.get("codes"), list) else [] state["codes"] = [item for item in legacy_codes if isinstance(item, dict)] + self._load_redeem_codes_from_session(session) state["ledger"] = self._load_shop_ledger_from_session(session) return state def _migrate_account_query_columns(self) -> None: with self._lock: session = self.Session() try: existing_columns = {column["name"] for column in inspect(self.engine).get_columns("accounts")} column_statements = { "owner_id": "ALTER TABLE accounts ADD COLUMN owner_id VARCHAR(255)", "shared": "ALTER TABLE accounts ADD COLUMN shared INTEGER NOT NULL DEFAULT 0", "status": "ALTER TABLE accounts ADD COLUMN status VARCHAR(64)", } for column_name, statement in column_statements.items(): if column_name not in existing_columns: try: session.execute(text(statement)) session.commit() except Exception as exc: session.rollback() refreshed_columns = { column["name"] for column in inspect(self.engine).get_columns("accounts") } if column_name not in refreshed_columns: print(f"[storage] account query column migration skipped: {exc}") return except Exception as exc: session.rollback() print(f"[storage] account query column migration skipped: {exc}") return finally: session.close() with self._lock: session = self.Session() try: for statement in ( "CREATE INDEX IF NOT EXISTS ix_accounts_owner_id ON accounts (owner_id)", "CREATE INDEX IF NOT EXISTS ix_accounts_shared ON accounts (shared)", "CREATE INDEX IF NOT EXISTS ix_accounts_status ON accounts (status)", ): try: session.execute(text(statement)) except Exception as exc: print(f"[storage] account query index migration skipped: {exc}") rows = session.query(AccountModel).all() for row in rows: item = self._account_from_row(row) if item is not None: self._apply_account_row(row, item) session.commit() except Exception as exc: session.rollback() print(f"[storage] account query metadata backfill skipped: {exc}") finally: session.close() def _migrate_legacy_redeem_codes(self) -> None: with self._lock: session = self.Session() try: self._acquire_transaction_lock(session, "shop_state") self._acquire_transaction_lock(session, "redeem_codes") state = self._load_shop_state_from_session(session, "default") legacy_codes = state.get("codes") if isinstance(state.get("codes"), list) else [] migrated = 0 for item in legacy_codes: if not isinstance(item, dict): continue code_id, code_hash = self._code_identity(item) if not code_id or not code_hash: continue row = session.query(RedeemCodeModel).filter_by(code_hash=code_hash).first() if row is not None: continue row = RedeemCodeModel(code_id=code_id, code_hash=code_hash, status="active", data="{}") self._apply_redeem_code_row(row, dict(item)) session.add(row) migrated += 1 if legacy_codes: state["codes"] = [ item for item in legacy_codes if not ( isinstance(item, dict) and str(item.get("id") or "").strip() and str(item.get("code_hash") or "").strip() ) ] self._save_raw_shop_state_to_session(session, "default", state) session.commit() if migrated: print(f"[storage] Migrated {migrated} redeem codes into indexed table") except Exception as e: session.rollback() print(f"[storage] redeem code migration skipped: {e}") finally: session.close() def _migrate_legacy_shop_ledger(self) -> None: with self._lock: session = self.Session() try: self._acquire_transaction_lock(session, "shop_state") self._acquire_transaction_lock(session, "shop_ledger") state = self._load_shop_state_from_session(session, "default") legacy_ledger = state.get("ledger") if isinstance(state.get("ledger"), list) else [] if not legacy_ledger: session.commit() return existing_ledger = self._load_shop_ledger_from_session(session) combined_ledger = [*existing_ledger, *legacy_ledger] self._sync_shop_ledger_to_session(session, combined_ledger) state["ledger"] = [] self._save_raw_shop_state_to_session(session, "default", state) session.commit() print(f"[storage] Migrated {len(legacy_ledger)} shop ledger entries into indexed table") except Exception as exc: session.rollback() print(f"[storage] shop ledger migration skipped: {exc}") finally: session.close() def load_redeem_codes(self) -> list[dict[str, Any]]: session = self.Session() try: return self._load_redeem_codes_from_session(session) finally: session.close() def load_shop_ledger(self) -> list[dict[str, Any]]: session = self.Session() try: return self._load_shop_ledger_from_session(session) finally: session.close() def append_shop_ledger(self, entries: list[dict[str, Any]]) -> None: """Append billing history without loading or rewriting prior entries.""" normalized_entries = [item for item in entries if isinstance(item, dict)] if not normalized_entries: return with self._lock: session = self.Session() try: self._acquire_transaction_lock(session, "shop_ledger") self._append_shop_ledger_to_session(session, normalized_entries) session.commit() except Exception as exc: session.rollback() raise exc finally: session.close() def get_redeem_code_by_hash(self, code_hash: str) -> dict[str, Any] | None: normalized_hash = str(code_hash or "").strip() if not normalized_hash: return None session = self.Session() try: row = session.query(RedeemCodeModel).filter_by(code_hash=normalized_hash).first() return self._redeem_code_from_row(row) if row is not None else None finally: session.close() def has_redeemed_code_in_batch(self, batch_id: str, user_id: str) -> bool: normalized_batch = str(batch_id or "").strip() normalized_user = str(user_id or "").strip() if not normalized_batch or not normalized_user: return False session = self.Session() try: return ( session.query(RedeemCodeModel) .filter_by(batch_id=normalized_batch, redeemed_by=normalized_user, status="redeemed") .first() is not None ) finally: session.close() def update_redeem_codes_locked( self, mutator: Callable[[dict[str, Any]], tuple[dict[str, Any], ResultT]], ) -> ResultT: """Mutate redeem codes and append ledger entries without touching shop_state.""" with self._lock: session = self.Session() try: self._acquire_transaction_lock(session, "redeem_codes") self._acquire_transaction_lock(session, "shop_ledger") state = { "codes": self._load_redeem_codes_from_session(session), "ledger": [], } next_state, result = mutator(state) next_state = next_state if isinstance(next_state, dict) else {} if "codes" not in next_state: next_state["codes"] = state.get("codes", []) next_codes = next_state.get("codes") if isinstance(next_state.get("codes"), list) else [] next_ledger = next_state.get("ledger") if isinstance(next_state.get("ledger"), list) else [] self._sync_redeem_codes_to_session(session, next_codes) self._append_shop_ledger_to_session(session, next_ledger) session.commit() return result except Exception as e: session.rollback() raise e finally: session.close() def delete_redeem_codes_locked( self, code_ids: list[str], mutator: Callable[[list[dict[str, Any]]], tuple[list[dict[str, Any]], ResultT]], ) -> ResultT: """Validate and delete selected redeem codes in one transaction.""" normalized_ids = list( dict.fromkeys( code_id for value in code_ids if (code_id := str(value or "").strip()) ) ) with self._lock: session = self.Session() try: self._acquire_transaction_lock(session, "redeem_codes") self._acquire_transaction_lock(session, "shop_ledger") rows: list[RedeemCodeModel] = [] for offset in range(0, len(normalized_ids), 500): chunk = normalized_ids[offset:offset + 500] query = session.query(RedeemCodeModel).filter(RedeemCodeModel.code_id.in_(chunk)) if self._is_postgresql: query = query.with_for_update() rows.extend(query.all()) selected_items: list[dict[str, Any]] = [] for row in rows: item = self._redeem_code_from_row(row) # Indexed columns are authoritative for identity and status. # This also keeps the redeemed guard safe if legacy JSON data # is stale or partially inconsistent. item["id"] = row.code_id item["status"] = row.status selected_items.append(item) ledger_entries, result = mutator(selected_items) for row in rows: session.delete(row) self._append_shop_ledger_to_session(session, ledger_entries) session.commit() return result except Exception: session.rollback() raise finally: session.close() def update_user_and_redeem_code_locked( self, *, code_hash: str, user_id: str, mutator: Callable[ [dict[str, Any] | None, dict[str, Any] | None, list[dict[str, Any]], dict[str, Any]], tuple[dict[str, Any] | None, dict[str, Any] | None, dict[str, Any], ResultT], ], ) -> ResultT: """Redeem one code atomically without reading or rewriting shop_state.""" normalized_hash = str(code_hash or "").strip() normalized_user = str(user_id or "").strip() with self._lock: session = self.Session() try: self._acquire_transaction_lock(session, "auth_keys") self._acquire_transaction_lock(session, "redeem_codes") self._acquire_transaction_lock(session, "shop_ledger") auth_row = session.query(AuthKeyModel).filter_by(key_id=normalized_user).first() auth_item = self._loads_dict(auth_row.data) if auth_row is not None else None code_row = session.query(RedeemCodeModel).filter_by(code_hash=normalized_hash).first() code_item = self._redeem_code_from_row(code_row) if code_row is not None else None batch_items: list[dict[str, Any]] = [] batch_id = str((code_item or {}).get("batch_id") or "").strip() if batch_id and normalized_user: batch_rows = ( session.query(RedeemCodeModel) .filter_by(batch_id=batch_id, redeemed_by=normalized_user, status="redeemed") .all() ) batch_items = [self._redeem_code_from_row(row) for row in batch_rows] state = {"codes": [], "ledger": []} next_auth_item, next_code_item, next_state, result = mutator(auth_item, code_item, batch_items, state) if next_auth_item is not None: key_id = str(next_auth_item.get("id") or "").strip() if not key_id: raise ValueError("用户不存在") payload = json.dumps(next_auth_item, ensure_ascii=False) if auth_row is None: session.add(AuthKeyModel(key_id=key_id, data=payload)) else: auth_row.data = payload if next_code_item is not None: next_code_id, next_code_hash = self._code_identity(next_code_item) if not next_code_id or not next_code_hash: raise ValueError("卡密不存在") if code_row is None: code_row = RedeemCodeModel( code_id=next_code_id, code_hash=next_code_hash, status=str(next_code_item.get("status") or "active").strip() or "active", data="{}", ) session.add(code_row) self._apply_redeem_code_row(code_row, next_code_item) normalized_next_state = next_state if isinstance(next_state, dict) else {} next_ledger = normalized_next_state.get("ledger") if isinstance(next_ledger, list): self._append_shop_ledger_to_session(session, next_ledger) session.commit() return result except Exception as e: session.rollback() raise e finally: session.close() def load_shop_state(self) -> dict[str, Any]: return self.load_named_state("default") def load_named_state(self, key: str) -> dict[str, Any]: session = self.Session() try: normalized_key = str(key or "").strip() or "default" if normalized_key == "default": return self._load_default_shop_state_from_session(session) return self._load_shop_state_from_session(session, normalized_key) finally: session.close() def save_shop_state(self, state: dict[str, Any]) -> None: self.save_named_state("default", state) def save_named_state(self, key: str, state: dict[str, Any]) -> None: with self._lock: session = self.Session() try: normalized_key = str(key or "").strip() or "default" if normalized_key == "default": if isinstance(state.get("codes"), list): self._sync_redeem_codes_to_session(session, state["codes"]) if isinstance(state.get("ledger"), list): self._sync_shop_ledger_to_session(session, state["ledger"]) self._save_shop_state_to_session(session, normalized_key, state) session.commit() except Exception as e: session.rollback() raise e finally: session.close() def update_shop_state_locked(self, mutator: Callable[[dict[str, Any]], tuple[dict[str, Any], ResultT]]) -> ResultT: """Load, mutate, and save shop state under one database transaction lock.""" with self._lock: session = self.Session() try: self._acquire_transaction_lock(session, "shop_state") self._acquire_transaction_lock(session, "redeem_codes") self._acquire_transaction_lock(session, "shop_ledger") state = self._load_default_shop_state_from_session(session) next_state, result = mutator(state) next_state = next_state if isinstance(next_state, dict) else {} if "codes" not in next_state: next_state["codes"] = state.get("codes", []) if "ledger" not in next_state: next_state["ledger"] = state.get("ledger", []) next_codes = next_state.get("codes") if isinstance(next_state.get("codes"), list) else [] next_ledger = next_state.get("ledger") if isinstance(next_state.get("ledger"), list) else [] self._sync_redeem_codes_to_session(session, next_codes) self._sync_shop_ledger_to_session(session, next_ledger) self._save_shop_state_to_session(session, "default", next_state) session.commit() return result except Exception as e: session.rollback() raise e finally: session.close() def update_user_and_shop_state_locked( self, mutator: Callable[[list[dict[str, Any]], dict[str, Any]], tuple[list[dict[str, Any]], dict[str, Any], ResultT]], ) -> ResultT: """Load, mutate, and save auth keys plus shop state in one transaction.""" with self._lock: session = self.Session() try: self._acquire_transaction_lock(session, "shop_state") self._acquire_transaction_lock(session, "auth_keys") self._acquire_transaction_lock(session, "redeem_codes") self._acquire_transaction_lock(session, "shop_ledger") auth_rows = session.query(AuthKeyModel).all() auth_items: list[dict[str, Any]] = [] for row in auth_rows: try: item_data = json.loads(row.data) except json.JSONDecodeError: continue if isinstance(item_data, dict): auth_items.append(item_data) shop_state = self._load_default_shop_state_from_session(session) next_auth_items, next_shop_state, result = mutator(auth_items, shop_state) next_shop_state = next_shop_state if isinstance(next_shop_state, dict) else {} if "codes" not in next_shop_state: next_shop_state["codes"] = shop_state.get("codes", []) if "ledger" not in next_shop_state: next_shop_state["ledger"] = shop_state.get("ledger", []) normalized_auth_items = [item for item in next_auth_items if isinstance(item, dict)] next_ids = {str(item.get("id") or "").strip() for item in normalized_auth_items} next_ids.discard("") for row in auth_rows: if str(row.key_id or "").strip() not in next_ids: session.delete(row) existing_by_id = {str(row.key_id or "").strip(): row for row in auth_rows} for item in normalized_auth_items: key_id = str(item.get("id") or "").strip() if not key_id: continue payload = json.dumps(item, ensure_ascii=False) row = existing_by_id.get(key_id) if row is None: session.add(AuthKeyModel(key_id=key_id, data=payload)) else: row.data = payload next_codes = next_shop_state.get("codes") if isinstance(next_shop_state.get("codes"), list) else [] self._sync_redeem_codes_to_session(session, next_codes) next_ledger = next_shop_state.get("ledger") if isinstance(next_shop_state.get("ledger"), list) else [] self._sync_shop_ledger_to_session(session, next_ledger) self._save_shop_state_to_session(session, "default", next_shop_state) session.commit() return result except Exception as e: session.rollback() raise e finally: session.close() def _load_rows(self, model: type[AccountModel] | type[AuthKeyModel]) -> list[dict[str, Any]]: session = self.Session() try: items = [] for row in session.query(model).all(): try: item_data = json.loads(row.data) if isinstance(item_data, dict): items.append(item_data) except json.JSONDecodeError: continue return items finally: session.close() def _save_rows( self, model: type[AccountModel] | type[AuthKeyModel], items: list[dict[str, Any]], source_key: str, target_key: str | None = None, ) -> None: session = self.Session() try: session.query(model).delete() for item in items: if not isinstance(item, dict): continue key_value = str(item.get(source_key) or "").strip() if not key_value: continue session.add( model( **{target_key or source_key: key_value}, data=json.dumps(item, ensure_ascii=False), ) ) session.commit() except Exception as e: session.rollback() raise e finally: session.close() @property def _is_postgresql(self) -> bool: normalized = self.database_url.lower() return normalized.startswith("postgresql://") or normalized.startswith("postgres://") @staticmethod def _lock_id(name: str) -> int: digest = hashlib.sha256(f"chatsam:{name}".encode("utf-8")).digest() return int.from_bytes(digest[:8], "big") & 0x7FFF_FFFF_FFFF_FFFF def _acquire_transaction_lock(self, session: Any, name: str) -> None: if not self._is_postgresql: return session.execute(text("SELECT pg_advisory_xact_lock(:lock_id)"), {"lock_id": self._lock_id(name)}) def health_check(self) -> dict[str, Any]: """健康检查""" try: session = self.Session() try: # 尝试执行简单查询 session.execute(text("SELECT 1")) result: dict[str, Any] = { "status": "healthy", "backend": "database", "database_url": self._mask_password(self.database_url), "table_scope": self.table_scope, } if self.table_scope in {"all", "accounts"}: result["account_count"] = session.query(AccountModel).count() if self.table_scope in {"all", "users"}: result["auth_key_count"] = session.query(AuthKeyModel).count() result["shop_state_count"] = session.query(ShopStateModel).count() result["redeem_code_count"] = session.query(RedeemCodeModel).count() result["shop_ledger_count"] = session.query(ShopLedgerModel).count() return result finally: session.close() except Exception as e: return { "status": "unhealthy", "backend": "database", "error": str(e), } def get_backend_info(self) -> dict[str, Any]: """获取存储后端信息""" db_type = "unknown" if "sqlite" in self.database_url: db_type = "sqlite" elif "postgresql" in self.database_url or "postgres" in self.database_url: db_type = "postgresql" elif "mysql" in self.database_url: db_type = "mysql" return { "type": "database", "db_type": db_type, "table_scope": self.table_scope, "description": f"数据库存储 ({db_type})", "database_url": self._mask_password(self.database_url), } @staticmethod def _mask_password(url: str) -> str: """隐藏数据库连接字符串中的密码""" if "://" not in url: return url try: protocol, rest = url.split("://", 1) if "@" in rest: credentials, host = rest.split("@", 1) if ":" in credentials: username, _ = credentials.split(":", 1) return f"{protocol}://{username}:****@{host}" return url except Exception: return url