from __future__ import annotations import os from pathlib import Path import shutil import sqlite3 import tempfile from threading import Event, RLock, Thread import time from typing import Iterable _SCHEMAS: dict[str, tuple[str, tuple[str, ...]]] = { "accounts": ( "CREATE TABLE IF NOT EXISTS accounts (id INTEGER PRIMARY KEY AUTOINCREMENT, access_token VARCHAR(2048) NOT NULL UNIQUE, owner_id VARCHAR(255), shared INTEGER NOT NULL DEFAULT 0, status VARCHAR(64), data TEXT NOT NULL)", ("id", "access_token", "owner_id", "shared", "status", "data"), ), "auth_keys": ( "CREATE TABLE IF NOT EXISTS auth_keys (id INTEGER PRIMARY KEY AUTOINCREMENT, key_id VARCHAR(255) NOT NULL UNIQUE, data TEXT NOT NULL)", ("id", "key_id", "data"), ), "shop_state": ( 'CREATE TABLE IF NOT EXISTS shop_state ("key" VARCHAR(64) PRIMARY KEY, data TEXT NOT NULL)', ("key", "data"), ), "redeem_codes": ( "CREATE TABLE IF NOT EXISTS redeem_codes (id INTEGER PRIMARY KEY AUTOINCREMENT, code_id VARCHAR(64) NOT NULL UNIQUE, code_hash VARCHAR(128) NOT NULL UNIQUE, batch_id VARCHAR(64), redeemed_by VARCHAR(255), status VARCHAR(32) NOT NULL, data TEXT NOT NULL)", ("id", "code_id", "code_hash", "batch_id", "redeemed_by", "status", "data"), ), "shop_ledger": ( "CREATE TABLE IF NOT EXISTS shop_ledger (id INTEGER PRIMARY KEY AUTOINCREMENT, ledger_id VARCHAR(64) NOT NULL UNIQUE, created_at VARCHAR(64), user_id VARCHAR(255), data TEXT NOT NULL)", ("id", "ledger_id", "created_at", "user_id", "data"), ), # 生图用量流水（data/credit_ledger.db）。它不在上面那几个库里，但同样跑在这套 # 镜像机制下；登记 schema 是为了数据库损坏时能被尽力抢救，而不是只能丢掉。 "credit_usage": ( "CREATE TABLE IF NOT EXISTS credit_usage (id INTEGER PRIMARY KEY AUTOINCREMENT, created_ts REAL NOT NULL, day TEXT NOT NULL, identity_id TEXT NOT NULL, role TEXT NOT NULL, model TEXT NOT NULL, status TEXT NOT NULL, image_count INTEGER NOT NULL DEFAULT 0, credits REAL NOT NULL DEFAULT 0, unit_price REAL NOT NULL DEFAULT 0, task_id TEXT NOT NULL DEFAULT '')", ("id", "created_ts", "day", "identity_id", "role", "model", "status", "image_count", "credits", "unit_price", "task_id"), ), } def _enabled() -> bool: return os.getenv("CHATGPT2API_SQLITE_BUCKET_SAFE_MODE", "true").strip().lower() not in {"0", "false", "no", "off"} def _integrity_ok(path: Path) -> bool: if not path.is_file() or path.stat().st_size == 0: return False connection: sqlite3.Connection | None = None try: connection = sqlite3.connect(f"file:{path.as_posix()}?mode=ro", uri=True, timeout=10) row = connection.execute("PRAGMA integrity_check").fetchone() return bool(row and str(row[0]).lower() == "ok") except sqlite3.Error: return False finally: if connection is not None: connection.close() def _copy_file(source: Path, destination: Path) -> None: destination.parent.mkdir(parents=True, exist_ok=True) temporary = destination.with_name(f".{destination.name}.copying-{os.getpid()}-{time.time_ns()}") try: shutil.copyfile(source, temporary) os.replace(temporary, destination) finally: temporary.unlink(missing_ok=True) class SQLiteMirrorService: """Run SQLite locally and persist only verified, immutable snapshots. Hugging Face Storage Buckets are mounted through a network filesystem. A live SQLite database performs page-level random writes and lock/journal coordination that such mounts cannot reliably preserve. This service keeps the live files on the container's local disk and treats the mounted bucket as snapshot storage only. """ def __init__(self) -> None: self._lock = RLock() self._runtime_dirs: dict[str, Path] = {} self._databases: dict[str, tuple[Path, Path]] = {} self._snapshot_signatures: dict[str, tuple[int, int]] = {} self._stop_event: Event | None = None self._thread: Thread | None = None @staticmethod def should_mirror(data_dir: Path) -> bool: if not _enabled(): return False persist_dir = Path(os.getenv("CHATGPT2API_PERSIST_DIR", "").strip() or "/app/persist") try: return persist_dir.exists() and data_dir.resolve(strict=False).is_relative_to(persist_dir.resolve(strict=False)) except (OSError, ValueError): return False def runtime_data_dir(self, data_dir: Path) -> Path: if not self.should_mirror(data_dir): return data_dir key = str(data_dir.resolve(strict=False)) with self._lock: existing = self._runtime_dirs.get(key) if existing is not None: return existing configured = os.getenv("CHATGPT2API_SQLITE_RUNTIME_DIR", "").strip() runtime = Path(configured) if configured else Path(tempfile.gettempdir()) / "chatgpt2api-sqlite" runtime.mkdir(parents=True, exist_ok=True) marker = data_dir / ".local_user_database_split_migrated" if marker.is_file(): _copy_file(marker, runtime / marker.name) for filename in ("accounts.db", "users.db"): self._prepare_database(data_dir, runtime, filename) self._runtime_dirs[key] = runtime return runtime def runtime_file(self, data_dir: Path, filename: str) -> Path: """Return a safe runtime path and register extra SQLite databases.""" runtime = self.runtime_data_dir(data_dir) if runtime == data_dir: return data_dir / filename target = runtime / filename if filename.endswith(".db"): with self._lock: if filename not in self._databases: self._prepare_database(data_dir, runtime, filename) return target source = data_dir / filename if source.is_file() and not target.exists(): _copy_file(source, target) return target def _snapshot_candidates(self, data_dir: Path, filename: str) -> Iterable[Path]: snapshot_dir = data_dir / "sqlite-snapshots" / filename.removesuffix(".db") if snapshot_dir.is_dir(): yield from sorted(snapshot_dir.glob("*.db"), reverse=True) legacy = data_dir / filename if legacy.is_file(): yield legacy def _prepare_database(self, data_dir: Path, runtime: Path, filename: str) -> None: live = runtime / filename if _integrity_ok(live): self._databases[filename] = (live, data_dir) return live.unlink(missing_ok=True) chosen: Path | None = None for candidate in self._snapshot_candidates(data_dir, filename): if _integrity_ok(candidate): _copy_file(candidate, live) chosen = candidate break if chosen is not None: print(f"[sqlite-mirror] restored {filename} from {chosen}") else: legacy = data_dir / filename if legacy.is_file(): chosen = self._salvage_database(legacy, live, data_dir) if chosen is None: live.parent.mkdir(parents=True, exist_ok=True) self._databases[filename] = (live, data_dir) def _salvage_database(self, source: Path, target: Path, data_dir: Path) -> Path | None: timestamp = time.strftime("%Y%m%d-%H%M%S", time.gmtime()) quarantine = data_dir / "sqlite-corrupt" / f"{source.stem}-{timestamp}.db" try: _copy_file(source, quarantine) except OSError as exc: print(f"[sqlite-mirror] failed to preserve corrupt {source.name}: {exc}") target.unlink(missing_ok=True) recovered: dict[str, int] = {} source_connection: sqlite3.Connection | None = None target_connection: sqlite3.Connection | None = None try: source_connection = sqlite3.connect(f"file:{source.as_posix()}?mode=ro", uri=True, timeout=10) target_connection = sqlite3.connect(str(target), timeout=30) target_connection.execute("PRAGMA journal_mode=DELETE") target_connection.execute("PRAGMA synchronous=FULL") table_names = {str(row[0]) for row in source_connection.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall()} for table_name, (schema, columns) in _SCHEMAS.items(): if table_name not in table_names: continue target_connection.execute(schema) quoted_columns = ",".join(f'"{column}"' for column in columns) placeholders = ",".join("?" for _ in columns) count = 0 try: cursor = source_connection.execute(f'SELECT {quoted_columns} FROM "{table_name}" ORDER BY rowid') except sqlite3.Error as exc: recovered[table_name] = 0 print(f"[sqlite-mirror] unable to start recovery for {table_name}: {exc}") continue while True: try: row = cursor.fetchone() except sqlite3.Error as exc: print(f"[sqlite-mirror] partial recovery for {table_name}: {exc}") break if row is None: break try: target_connection.execute( f'INSERT OR IGNORE INTO "{table_name}" ({quoted_columns}) VALUES ({placeholders})', row, ) count += 1 except sqlite3.Error: continue recovered[table_name] = count target_connection.commit() row = target_connection.execute("PRAGMA integrity_check").fetchone() if not row or str(row[0]).lower() != "ok": raise sqlite3.DatabaseError("recovered database failed integrity_check") print(f"[sqlite-mirror] salvaged corrupt {source.name}: {recovered}; original={quarantine}") return quarantine except sqlite3.Error as exc: target.unlink(missing_ok=True) print(f"[sqlite-mirror] automatic salvage failed for {source.name}: {exc}; original={quarantine}") return None finally: if target_connection is not None: target_connection.close() if source_connection is not None: source_connection.close() def _snapshot_database(self, filename: str, live: Path, data_dir: Path) -> Path | None: if not live.is_file(): return None snapshot_dir = data_dir / "sqlite-snapshots" / filename.removesuffix(".db") snapshot_dir.mkdir(parents=True, exist_ok=True) timestamp = time.strftime("%Y%m%d-%H%M%S", time.gmtime()) local_copy = live.parent / f".{filename}.snapshot-{os.getpid()}-{time.time_ns()}" source_connection: sqlite3.Connection | None = None target_connection: sqlite3.Connection | None = None try: source_connection = sqlite3.connect(str(live), timeout=30) target_connection = sqlite3.connect(str(local_copy), timeout=30) source_connection.backup(target_connection) target_connection.close() target_connection = None if not _integrity_ok(local_copy): raise sqlite3.DatabaseError("snapshot integrity_check failed") destination = snapshot_dir / f"{timestamp}-{time.time_ns()}.db" _copy_file(local_copy, destination) if not _integrity_ok(destination): destination.unlink(missing_ok=True) raise sqlite3.DatabaseError("uploaded snapshot integrity_check failed") snapshots = sorted(snapshot_dir.glob("*.db"), reverse=True) keep = max(2, int(os.getenv("CHATGPT2API_SQLITE_SNAPSHOT_KEEP", "5"))) for stale in snapshots[keep:]: stale.unlink(missing_ok=True) return destination except (OSError, sqlite3.Error) as exc: print(f"[sqlite-mirror] snapshot failed for {filename}: {exc}") return None finally: if target_connection is not None: target_connection.close() if source_connection is not None: source_connection.close() local_copy.unlink(missing_ok=True) def snapshot_all(self, *, force: bool = False) -> list[str]: completed: list[str] = [] with self._lock: databases = list(self._databases.items()) for filename, (live, data_dir) in databases: try: stat = live.stat() signature = (stat.st_mtime_ns, stat.st_size) except OSError: continue if not force and self._snapshot_signatures.get(filename) == signature: continue if self._snapshot_database(filename, live, data_dir) is not None: completed.append(filename) self._snapshot_signatures[filename] = signature if completed: print(f"[sqlite-mirror] verified snapshots saved: {', '.join(completed)}") return completed def start(self) -> None: with self._lock: if self._thread is not None or not self._databases: return self.snapshot_all(force=True) stop_event = Event() interval = max(10, int(os.getenv("CHATGPT2API_SQLITE_SNAPSHOT_INTERVAL_SECS", "30"))) def worker() -> None: while not stop_event.wait(interval): self.snapshot_all() self._stop_event = stop_event self._thread = Thread(target=worker, name="sqlite-bucket-snapshot", daemon=True) self._thread.start() def stop(self) -> None: with self._lock: stop_event = self._stop_event thread = self._thread self._stop_event = None self._thread = None if stop_event is None: return stop_event.set() if thread is not None: thread.join(timeout=2) self.snapshot_all(force=True) sqlite_mirror_service = SQLiteMirrorService()