services/sqlite_mirror_service.py

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