| from __future__ import annotations |
|
|
| import os |
| import sqlite3 |
| import tempfile |
| from pathlib import Path |
| from threading import RLock |
| from typing import Any |
|
|
| from services.storage.account_pool_storage import AccountPoolStorageBackend, migrate_personal_accounts_to_shared_once |
| from services.storage.base import StorageBackend |
| from services.storage.database_storage import DatabaseStorageBackend |
| from services.storage.git_storage import GitStorageBackend |
| from services.storage.json_storage import JSONStorageBackend |
| from services.sqlite_mirror_service import sqlite_mirror_service |
|
|
|
|
| def _default_backend_type(data_dir: Path) -> str: |
| if os.getenv("STORAGE_BACKEND") is not None: |
| return os.getenv("STORAGE_BACKEND", "json").lower().strip() |
| persist_dir = Path(os.getenv("CHATGPT2API_PERSIST_DIR", "").strip() or "/app/persist") |
| if persist_dir.exists() or str(data_dir).startswith(str(persist_dir)): |
| return "sqlite" |
| return "json" |
|
|
|
|
| def _sqlite_database_url(data_dir: Path) -> str: |
| return f"sqlite:///{data_dir / 'accounts.db'}" |
|
|
|
|
| def _sqlite_user_database_url(data_dir: Path) -> str: |
| return f"sqlite:///{data_dir / 'users.db'}" |
|
|
|
|
| def _is_local_accounts_database_url(database_url: str, data_dir: Path) -> bool: |
| if not database_url.startswith("sqlite:///"): |
| return False |
| configured_path = Path(database_url.removeprefix("sqlite:///")) |
| try: |
| return configured_path.resolve(strict=False) == (data_dir / "accounts.db").resolve(strict=False) |
| except OSError: |
| return str(configured_path) == str(data_dir / "accounts.db") |
|
|
|
|
| _LOCAL_SQLITE_SPLIT_LOCK = RLock() |
|
|
|
|
| def _local_sqlite_split_enabled() -> bool: |
| return os.getenv("CHATGPT2API_SPLIT_LOCAL_DATABASES", "true").strip().lower() not in {"0", "false", "no", "off"} |
|
|
|
|
| def _sqlite_values(connection: sqlite3.Connection, query: str) -> set[object]: |
| return {row[0] for row in connection.execute(query).fetchall()} |
|
|
|
|
| def _migrate_local_sqlite_users_once(data_dir: Path) -> None: |
| """Move user/shop tables from accounts.db to users.db, then remove the legacy copies. |
| |
| The migration is idempotent. Target rows win on natural-key conflicts so a |
| partially migrated deployment never overwrites newer user balances or codes. |
| """ |
| marker_path = data_dir / ".local_user_database_split_migrated" |
| accounts_path = data_dir / "accounts.db" |
| users_path = data_dir / "users.db" |
| backup_path = data_dir / "accounts.before-user-split.db" |
| with _LOCAL_SQLITE_SPLIT_LOCK: |
| if marker_path.exists() and users_path.exists(): |
| return |
|
|
| |
| legacy_backend = DatabaseStorageBackend(f"sqlite:///{accounts_path}", table_scope="all") |
| _migrate_git_to_database_if_empty(legacy_backend) |
| user_backend = DatabaseStorageBackend(f"sqlite:///{users_path}", table_scope="users") |
| legacy_backend.engine.dispose() |
| user_backend.engine.dispose() |
|
|
| source = sqlite3.connect(str(accounts_path), timeout=30) |
| target = sqlite3.connect(str(users_path), timeout=30) |
| try: |
| for connection in (source, target): |
| connection.execute("PRAGMA busy_timeout=30000") |
| connection.execute("PRAGMA synchronous=FULL") |
| connection.execute("PRAGMA journal_mode=DELETE") |
|
|
| if not backup_path.exists(): |
| backup = sqlite3.connect(str(backup_path)) |
| try: |
| source.backup(backup) |
| finally: |
| backup.close() |
|
|
| target.execute("BEGIN IMMEDIATE") |
| target.executemany( |
| "INSERT OR IGNORE INTO auth_keys(key_id,data) VALUES(?,?)", |
| source.execute("SELECT key_id,data FROM auth_keys ORDER BY id").fetchall(), |
| ) |
| target.executemany( |
| 'INSERT OR IGNORE INTO shop_state("key",data) VALUES(?,?)', |
| source.execute('SELECT "key",data FROM shop_state ORDER BY "key"').fetchall(), |
| ) |
| target.executemany( |
| "INSERT OR IGNORE INTO redeem_codes(code_id,code_hash,batch_id,redeemed_by,status,data) VALUES(?,?,?,?,?,?)", |
| source.execute( |
| "SELECT code_id,code_hash,batch_id,redeemed_by,status,data FROM redeem_codes ORDER BY id" |
| ).fetchall(), |
| ) |
| target.executemany( |
| "INSERT OR IGNORE INTO shop_ledger(ledger_id,created_at,user_id,data) VALUES(?,?,?,?)", |
| source.execute("SELECT ledger_id,created_at,user_id,data FROM shop_ledger ORDER BY id").fetchall(), |
| ) |
| target.commit() |
|
|
| checks = [ |
| ( |
| _sqlite_values(source, "SELECT key_id FROM auth_keys"), |
| _sqlite_values(target, "SELECT key_id FROM auth_keys"), |
| "auth_keys", |
| ), |
| ( |
| _sqlite_values(source, 'SELECT "key" FROM shop_state'), |
| _sqlite_values(target, 'SELECT "key" FROM shop_state'), |
| "shop_state", |
| ), |
| ( |
| _sqlite_values(source, "SELECT code_hash FROM redeem_codes"), |
| _sqlite_values(target, "SELECT code_hash FROM redeem_codes"), |
| "redeem_codes", |
| ), |
| ( |
| _sqlite_values(source, "SELECT ledger_id FROM shop_ledger"), |
| _sqlite_values(target, "SELECT ledger_id FROM shop_ledger"), |
| "shop_ledger", |
| ), |
| ] |
| for source_keys, target_keys, table_name in checks: |
| missing = source_keys - target_keys |
| if missing: |
| raise RuntimeError(f"local user database split verification failed for {table_name}: {len(missing)} rows missing") |
| if target.execute("PRAGMA integrity_check").fetchone()[0] != "ok": |
| raise RuntimeError("users.db failed integrity_check after migration") |
|
|
| source.execute("BEGIN IMMEDIATE") |
| for table_name in ("shop_ledger", "redeem_codes", "shop_state", "auth_keys"): |
| source.execute(f'DROP TABLE IF EXISTS "{table_name}"') |
| source.commit() |
| source.execute("ANALYZE") |
| source.execute("VACUUM") |
| if source.execute("PRAGMA integrity_check").fetchone()[0] != "ok": |
| raise RuntimeError("accounts.db failed integrity_check after user table removal") |
|
|
| marker_path.write_text("completed\n", encoding="utf-8") |
| print( |
| "[storage] Split local SQLite storage: " |
| f"accounts={accounts_path}, users={users_path}, backup={backup_path}" |
| ) |
| except Exception: |
| target.rollback() |
| source.rollback() |
| raise |
| finally: |
| target.close() |
| source.close() |
|
|
|
|
| def _is_database_empty(backend: DatabaseStorageBackend) -> bool: |
| return not backend.load_accounts() and not backend.load_auth_keys() and not backend.load_shop_state() |
|
|
|
|
| def _shop_state_has_data(state: object) -> bool: |
| if not isinstance(state, dict): |
| return False |
| return bool(state.get("codes") or state.get("ledger")) |
|
|
|
|
| def _is_user_storage_empty(backend: StorageBackend) -> bool: |
| return not backend.load_auth_keys() and not _shop_state_has_data(backend.load_shop_state()) |
|
|
|
|
| def _merge_auth_keys_preserving_target(source_items: list[dict[str, Any]], target_items: list[dict[str, Any]]) -> list[dict[str, Any]]: |
| merged: dict[str, dict[str, Any]] = {} |
| for item in source_items: |
| item_id = str(item.get("id") or "").strip() |
| if item_id: |
| merged[item_id] = dict(item) |
| for item in target_items: |
| item_id = str(item.get("id") or "").strip() |
| if item_id: |
| merged[item_id] = dict(item) |
| return list(merged.values()) |
|
|
|
|
| def _merge_shop_state_preserving_target(source_state: dict[str, Any], target_state: dict[str, Any]) -> dict[str, Any]: |
| def merge_items(key: str, identity_keys: tuple[str, ...]) -> list[dict[str, Any]]: |
| merged: dict[str, dict[str, Any]] = {} |
| source_items = source_state.get(key) if isinstance(source_state.get(key), list) else [] |
| target_items = target_state.get(key) if isinstance(target_state.get(key), list) else [] |
| for item in [*source_items, *target_items]: |
| if not isinstance(item, dict): |
| continue |
| item_id = "" |
| for identity_key in identity_keys: |
| item_id = str(item.get(identity_key) or "").strip() |
| if item_id: |
| break |
| if item_id: |
| merged[item_id] = dict(item) |
| return list(merged.values()) |
|
|
| return { |
| "codes": merge_items("codes", ("id", "code_hash")), |
| "ledger": sorted(merge_items("ledger", ("id",)), key=lambda item: str(item.get("time") or "")), |
| } |
|
|
|
|
| def _migrate_local_users_to_shared_once(source: StorageBackend, target: StorageBackend, marker_path: Path) -> None: |
| if source is target: |
| return |
| if os.getenv("CHATGPT2API_AUTO_MIGRATE_USER_STORAGE", "").strip().lower() not in {"1", "true", "yes", "on"}: |
| return |
| if os.getenv("CHATGPT2API_SKIP_USER_STORAGE_MIGRATION", "").strip().lower() in {"1", "true", "yes", "on"}: |
| return |
| if marker_path.exists(): |
| return |
|
|
| try: |
| source_auth_keys = source.load_auth_keys() |
| source_shop_state = source.load_shop_state() |
| except Exception as exc: |
| print(f"[storage] shared user migration skipped, source load failed: {exc}") |
| return |
|
|
| if not source_auth_keys and not _shop_state_has_data(source_shop_state): |
| return |
|
|
| def merge(target_auth_keys: list[dict[str, Any]], target_shop_state: dict[str, Any]) -> tuple[list[dict[str, Any]], dict[str, Any]]: |
| return ( |
| _merge_auth_keys_preserving_target(source_auth_keys, target_auth_keys), |
| _merge_shop_state_preserving_target( |
| source_shop_state if isinstance(source_shop_state, dict) else {}, |
| target_shop_state if isinstance(target_shop_state, dict) else {}, |
| ), |
| ) |
|
|
| try: |
| update_user_and_shop_state_locked = getattr(target, "update_user_and_shop_state_locked", None) |
| if callable(update_user_and_shop_state_locked): |
| def mutate( |
| target_auth_keys: list[dict[str, Any]], |
| target_shop_state: dict[str, Any], |
| ) -> tuple[list[dict[str, Any]], dict[str, Any], dict[str, int]]: |
| merged_auth_keys, merged_shop_state = merge(target_auth_keys, target_shop_state) |
| result = { |
| "source_auth_keys": len(source_auth_keys), |
| "target_auth_keys": len(target_auth_keys), |
| "merged_auth_keys": len(merged_auth_keys), |
| "source_codes": len(source_shop_state.get("codes") or []) if isinstance(source_shop_state, dict) else 0, |
| "merged_codes": len(merged_shop_state.get("codes") or []), |
| } |
| return merged_auth_keys, merged_shop_state, result |
|
|
| result = update_user_and_shop_state_locked(mutate) |
| else: |
| target_auth_keys = target.load_auth_keys() |
| target_shop_state = target.load_shop_state() |
| merged_auth_keys, merged_shop_state = merge(target_auth_keys, target_shop_state) |
| target.save_auth_keys(merged_auth_keys) |
| target.save_shop_state(merged_shop_state) |
| result = { |
| "source_auth_keys": len(source_auth_keys), |
| "target_auth_keys": len(target_auth_keys), |
| "merged_auth_keys": len(merged_auth_keys), |
| "source_codes": len(source_shop_state.get("codes") or []) if isinstance(source_shop_state, dict) else 0, |
| "merged_codes": len(merged_shop_state.get("codes") or []), |
| } |
| marker_path.parent.mkdir(parents=True, exist_ok=True) |
| marker_path.write_text("completed\n", encoding="utf-8") |
| print( |
| "[storage] Merged local user data to shared database: " |
| f"source_auth_keys={result['source_auth_keys']}, " |
| f"target_auth_keys={result['target_auth_keys']}, " |
| f"merged_auth_keys={result['merged_auth_keys']}, " |
| f"source_codes={result['source_codes']}, " |
| f"merged_codes={result['merged_codes']}" |
| ) |
| except Exception as exc: |
| print(f"[storage] shared user migration failed: {exc}") |
|
|
|
|
| def _build_git_backend_from_env() -> GitStorageBackend | None: |
| repo_url = os.getenv("GIT_REPO_URL", "").strip() |
| token = os.getenv("GIT_TOKEN", "").strip() |
| if not repo_url: |
| return None |
| branch = os.getenv("GIT_BRANCH", "main").strip() |
| file_path = os.getenv("GIT_FILE_PATH", "accounts.json").strip() |
| auth_keys_file_path = os.getenv("GIT_AUTH_KEYS_FILE_PATH", "auth_keys.json").strip() |
| shop_state_file_path = os.getenv("GIT_SHOP_STATE_FILE_PATH", "shop_state.json").strip() |
| cache_dir = Path(os.getenv("GIT_CACHE_DIR", "").strip() or (Path(tempfile.gettempdir()) / "chatgpt2api_git_cache")) |
| return GitStorageBackend( |
| repo_url=repo_url, |
| token=token, |
| branch=branch, |
| file_path=file_path, |
| auth_keys_file_path=auth_keys_file_path, |
| shop_state_file_path=shop_state_file_path, |
| local_cache_dir=cache_dir, |
| ) |
|
|
|
|
| def _migrate_git_to_database_if_empty(backend: DatabaseStorageBackend) -> None: |
| if os.getenv("CHATGPT2API_SKIP_STORAGE_MIGRATION", "").strip().lower() in {"1", "true", "yes", "on"}: |
| return |
| try: |
| if not _is_database_empty(backend): |
| return |
| except Exception as exc: |
| print(f"[storage] skip migration, database check failed: {exc}") |
| return |
|
|
| git_backend = _build_git_backend_from_env() |
| if git_backend is None: |
| return |
|
|
| try: |
| accounts = git_backend.load_accounts() |
| auth_keys = git_backend.load_auth_keys() |
| shop_state = git_backend.load_shop_state() |
| except Exception as exc: |
| print(f"[storage] Git to SQLite migration skipped: {exc}") |
| return |
|
|
| if not accounts and not auth_keys and not shop_state: |
| return |
|
|
| print( |
| "[storage] Migrating Git data to SQLite: " |
| f"{len(accounts)} accounts, {len(auth_keys)} auth keys, shop_state={bool(shop_state)}" |
| ) |
| if accounts: |
| backend.save_accounts(accounts) |
| if auth_keys: |
| backend.save_auth_keys(auth_keys) |
| if shop_state: |
| backend.save_shop_state(shop_state) |
|
|
|
|
| def create_storage_backend(data_dir: Path) -> StorageBackend: |
| """ |
| 根据环境变量创建存储后端 |
| |
| 环境变量： |
| - STORAGE_BACKEND: json|sqlite|postgres|git (默认 json) |
| - DATABASE_URL: 数据库连接字符串 (用于 sqlite/postgres) |
| - GIT_REPO_URL: Git 仓库地址 (用于 git) |
| - GIT_TOKEN: Git 访问令牌 (用于 git) |
| - GIT_BRANCH: Git 分支 (默认 main) |
| - GIT_FILE_PATH: Git 仓库中的文件路径 (默认 accounts.json) |
| - GIT_CACHE_DIR: Git 本地缓存目录 (默认系统临时目录) |
| """ |
| backend_type = _default_backend_type(data_dir) |
| |
| print(f"[storage] Initializing storage backend: {backend_type}") |
| |
| if backend_type == "json": |
| |
| file_path = data_dir / "accounts.json" |
| auth_keys_path = data_dir / "auth_keys.json" |
| shop_state_path = data_dir / "shop_state.json" |
| print(f"[storage] Using JSON storage: {file_path}") |
| return JSONStorageBackend(file_path, auth_keys_path, shop_state_path) |
| |
| elif backend_type in ("sqlite", "postgres", "postgresql", "mysql", "database"): |
| |
| database_url = os.getenv("DATABASE_URL", "").strip() |
| sqlite_data_dir = data_dir |
| |
| if not database_url: |
| |
| sqlite_data_dir = sqlite_mirror_service.runtime_data_dir(data_dir) |
| database_url = _sqlite_database_url(sqlite_data_dir) |
| print(f"[storage] No DATABASE_URL provided, using local SQLite: {database_url}") |
| else: |
| if _is_local_accounts_database_url(database_url, data_dir): |
| sqlite_data_dir = sqlite_mirror_service.runtime_data_dir(data_dir) |
| database_url = _sqlite_database_url(sqlite_data_dir) |
| print(f"[storage] Using database storage: {_mask_password(database_url)}") |
|
|
| if ( |
| backend_type in {"sqlite", "database"} |
| and _local_sqlite_split_enabled() |
| and _is_local_accounts_database_url(database_url, sqlite_data_dir) |
| ): |
| _migrate_local_sqlite_users_once(sqlite_data_dir) |
| return DatabaseStorageBackend(database_url, table_scope="accounts") |
|
|
| backend = DatabaseStorageBackend(database_url) |
| if backend_type in {"sqlite", "database"}: |
| _migrate_git_to_database_if_empty(backend) |
| return backend |
| |
| elif backend_type == "git": |
| |
| repo_url = os.getenv("GIT_REPO_URL", "").strip() |
| token = os.getenv("GIT_TOKEN", "").strip() |
| branch = os.getenv("GIT_BRANCH", "main").strip() |
| file_path = os.getenv("GIT_FILE_PATH", "accounts.json").strip() |
| auth_keys_file_path = os.getenv("GIT_AUTH_KEYS_FILE_PATH", "auth_keys.json").strip() |
| shop_state_file_path = os.getenv("GIT_SHOP_STATE_FILE_PATH", "shop_state.json").strip() |
| |
| if not repo_url: |
| raise ValueError( |
| "GIT_REPO_URL is required when using git storage backend. " |
| "Please set GIT_REPO_URL environment variable." |
| ) |
| |
| print(f"[storage] Using Git storage: {_mask_token(repo_url)}, branch: {branch}, file: {file_path}") |
| |
| cache_dir = Path(os.getenv("GIT_CACHE_DIR", "").strip() or (Path(tempfile.gettempdir()) / "chatgpt2api_git_cache")) |
| return GitStorageBackend( |
| repo_url=repo_url, |
| token=token, |
| branch=branch, |
| file_path=file_path, |
| auth_keys_file_path=auth_keys_file_path, |
| shop_state_file_path=shop_state_file_path, |
| local_cache_dir=cache_dir, |
| ) |
| |
| else: |
| raise ValueError( |
| f"Unknown storage backend: {backend_type}. " |
| f"Supported backends: json, sqlite, postgres, git" |
| ) |
|
|
|
|
| def create_user_storage_backend(data_dir: Path, fallback_backend: StorageBackend | None = None) -> StorageBackend: |
| """Create the optional shared backend for users, credits, WeChat bindings, and redeem codes. |
| |
| A configured SHARED_USER_DATABASE_URL still takes precedence. Otherwise the default |
| local SQLite deployment uses users.db, isolated from the account-pool accounts.db. |
| """ |
| database_url = os.getenv("SHARED_USER_DATABASE_URL", "").strip() |
| if not database_url: |
| sqlite_data_dir = sqlite_mirror_service.runtime_data_dir(data_dir) |
| if _local_sqlite_split_enabled() and isinstance(fallback_backend, DatabaseStorageBackend): |
| if fallback_backend.table_scope == "accounts" and _is_local_accounts_database_url( |
| fallback_backend.database_url, sqlite_data_dir |
| ): |
| _migrate_local_sqlite_users_once(sqlite_data_dir) |
| user_database_url = _sqlite_user_database_url(sqlite_data_dir) |
| print(f"[storage] Using isolated local user SQLite: {user_database_url}") |
| return DatabaseStorageBackend(user_database_url, table_scope="users") |
| return fallback_backend or create_storage_backend(data_dir) |
|
|
| print(f"[storage] Using shared user database storage: {_mask_password(database_url)}") |
| backend = DatabaseStorageBackend(database_url) |
| if fallback_backend is not None: |
| migration_source = fallback_backend |
| local_user_backend: DatabaseStorageBackend | None = None |
| if ( |
| isinstance(fallback_backend, DatabaseStorageBackend) |
| and fallback_backend.table_scope == "accounts" |
| and _is_local_accounts_database_url(fallback_backend.database_url, data_dir) |
| ): |
| local_user_backend = DatabaseStorageBackend(_sqlite_user_database_url(data_dir), table_scope="users") |
| migration_source = local_user_backend |
| try: |
| _migrate_local_users_to_shared_once( |
| migration_source, |
| backend, |
| data_dir / ".shared_user_storage_migrated", |
| ) |
| finally: |
| if local_user_backend is not None: |
| local_user_backend.engine.dispose() |
| return backend |
|
|
|
|
| def create_account_pool_storage_backend( |
| data_dir: Path, |
| site_storage: StorageBackend, |
| user_storage: StorageBackend, |
| ) -> StorageBackend: |
| if not os.getenv("SHARED_USER_DATABASE_URL", "").strip() or site_storage is user_storage: |
| return site_storage |
| migrate_personal_accounts_to_shared_once( |
| site_storage, |
| user_storage, |
| data_dir / ".shared_account_pool_storage_migrated", |
| ) |
| return AccountPoolStorageBackend(site_storage, user_storage) |
|
|
|
|
| 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 |
|
|
|
|
| def _mask_token(url: str) -> str: |
| """隐藏 URL 中的 token""" |
| if "@" in url and "://" in url: |
| protocol, rest = url.split("://", 1) |
| if "@" in rest: |
| _, host = rest.split("@", 1) |
| return f"{protocol}://****@{host}" |
| return url |
|
|
