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, 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 DatabaseStorageBackend(StorageBackend): """数据库存储后端(支持 SQLite、PostgreSQL、MySQL 等)""" def __init__(self, database_url: str): self.database_url = database_url 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 self.engine = create_engine( database_url, pool_pre_ping=True, # 自动检测连接是否有效 pool_recycle=3600, # 1小时回收连接 connect_args=connect_args, ) Base.metadata.create_all(self.engine) self.Session = sessionmaker(bind=self.engine) self._lock = RLock() self._migrate_account_query_columns() self._migrate_legacy_redeem_codes() 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)) else: 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()) ] 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 payload = json.dumps(payload_state if isinstance(payload_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) row.code_id = code_id row.code_hash = code_hash row.batch_id = str(item.get("batch_id") or "").strip() or None row.redeemed_by = str(item.get("redeemed_by") or "").strip() or None row.status = str(item.get("status") or "active").strip() or "active" row.data = json.dumps(item, ensure_ascii=False) 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) 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: self._save_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 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: state = self._load_shop_state_from_session(session, "default") ledger = state.get("ledger") if isinstance(state.get("ledger"), list) else [] return [item for item in ledger if isinstance(item, dict)] 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: """Load, mutate, and save redeem codes plus shop ledger in one transaction.""" 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") codes = self._load_redeem_codes_from_session(session) legacy_codes = state.get("codes") if isinstance(state.get("codes"), list) else [] state["codes"] = [item for item in legacy_codes if isinstance(item, dict)] + codes 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 [] self._sync_redeem_codes_to_session(session, next_codes) 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_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: """Update one user and one redeem code without loading every code.""" 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, "shop_state") self._acquire_transaction_lock(session, "auth_keys") self._acquire_transaction_lock(session, "redeem_codes") 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 = self._load_shop_state_from_session(session, "default") 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) self._save_shop_state_to_session(session, "default", next_state if isinstance(next_state, dict) else {}) 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" state = self._load_shop_state_from_session(session, normalized_key) if normalized_key == "default": codes = self._load_redeem_codes_from_session(session) if codes: legacy_codes = state.get("codes") if isinstance(state.get("codes"), list) else [] state["codes"] = [item for item in legacy_codes if isinstance(item, dict)] + codes return state 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" and isinstance(state.get("codes"), list): self._sync_redeem_codes_to_session(session, state["codes"]) 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") state = self._load_shop_state_from_session(session, "default") codes = self._load_redeem_codes_from_session(session) if codes: legacy_codes = state.get("codes") if isinstance(state.get("codes"), list) else [] state["codes"] = [item for item in legacy_codes if isinstance(item, dict)] + codes 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 [] self._sync_redeem_codes_to_session(session, next_codes) 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") 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_shop_state_from_session(session, "default") codes = self._load_redeem_codes_from_session(session) if codes: legacy_codes = shop_state.get("codes") if isinstance(shop_state.get("codes"), list) else [] shop_state["codes"] = [item for item in legacy_codes if isinstance(item, dict)] + codes 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", []) 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) 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")) count = session.query(AccountModel).count() auth_key_count = session.query(AuthKeyModel).count() shop_state_count = session.query(ShopStateModel).count() redeem_code_count = session.query(RedeemCodeModel).count() return { "status": "healthy", "backend": "database", "database_url": self._mask_password(self.database_url), "account_count": count, "auth_key_count": auth_key_count, "shop_state_count": shop_state_count, "redeem_code_count": redeem_code_count, } 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, "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