| 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"), |
| ), |
| |
| |
| "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() |
|
|
