ComfyUI-LM-Remote/stats_db.py

"""SQLite storage for CivitAI extended stats."""
from __future__ import annotations

import sqlite3
import time
from pathlib import Path


_PACKAGE_DIR = Path(__file__).resolve().parent
_DEFAULT_DB_PATH = _PACKAGE_DIR / "civitai_stats.db"

_SCHEMA = """
CREATE TABLE IF NOT EXISTS model_stats (
    sha256 TEXT PRIMARY KEY,
    civitai_model_id INTEGER,
    civitai_version_id INTEGER,
    download_count INTEGER DEFAULT 0,
    rating REAL DEFAULT 0,
    rating_count INTEGER DEFAULT 0,
    thumbs_up_count INTEGER DEFAULT 0,
    fetched_at REAL
);
"""


class StatsDB:
    """Thin wrapper around a SQLite database for CivitAI stats."""

    def __init__(self, db_path: Path | None = None):
        self._db_path = db_path or _DEFAULT_DB_PATH
        self._conn: sqlite3.Connection | None = None

    def init(self) -> None:
        """Create the database and table if they don't exist."""
        self._conn = sqlite3.connect(str(self._db_path), check_same_thread=False)
        self._conn.row_factory = sqlite3.Row
        self._conn.execute("PRAGMA journal_mode=WAL")
        self._conn.executescript(_SCHEMA)

    def _ensure_conn(self) -> sqlite3.Connection:
        if self._conn is None:
            self.init()
        return self._conn  # type: ignore[return-value]

    def upsert(self, sha256: str, data: dict) -> None:
        """Insert or update stats for a single model."""
        conn = self._ensure_conn()
        conn.execute(
            """INSERT INTO model_stats
               (sha256, civitai_model_id, civitai_version_id,
                download_count, rating, rating_count, thumbs_up_count, fetched_at)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?)
               ON CONFLICT(sha256) DO UPDATE SET
                 civitai_model_id = COALESCE(excluded.civitai_model_id, civitai_model_id),
                 civitai_version_id = COALESCE(excluded.civitai_version_id, civitai_version_id),
                 download_count = excluded.download_count,
                 rating = excluded.rating,
                 rating_count = excluded.rating_count,
                 thumbs_up_count = excluded.thumbs_up_count,
                 fetched_at = excluded.fetched_at
            """,
            (
                sha256,
                data.get("civitai_model_id"),
                data.get("civitai_version_id"),
                data.get("download_count", 0),
                data.get("rating", 0),
                data.get("rating_count", 0),
                data.get("thumbs_up_count", 0),
                time.time(),
            ),
        )
        conn.commit()

    def upsert_batch(self, rows: list[tuple[str, dict]]) -> None:
        """Insert or update stats for multiple models."""
        if not rows:
            return
        conn = self._ensure_conn()
        now = time.time()
        conn.executemany(
            """INSERT INTO model_stats
               (sha256, civitai_model_id, civitai_version_id,
                download_count, rating, rating_count, thumbs_up_count, fetched_at)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?)
               ON CONFLICT(sha256) DO UPDATE SET
                 civitai_model_id = COALESCE(excluded.civitai_model_id, civitai_model_id),
                 civitai_version_id = COALESCE(excluded.civitai_version_id, civitai_version_id),
                 download_count = excluded.download_count,
                 rating = excluded.rating,
                 rating_count = excluded.rating_count,
                 thumbs_up_count = excluded.thumbs_up_count,
                 fetched_at = excluded.fetched_at
            """,
            [
                (
                    sha256,
                    d.get("civitai_model_id"),
                    d.get("civitai_version_id"),
                    d.get("download_count", 0),
                    d.get("rating", 0),
                    d.get("rating_count", 0),
                    d.get("thumbs_up_count", 0),
                    now,
                )
                for sha256, d in rows
            ],
        )
        conn.commit()

    def get_by_hashes(self, hashes: list[str]) -> dict[str, dict]:
        """Return stats keyed by sha256 for the given hashes."""
        if not hashes:
            return {}
        conn = self._ensure_conn()
        placeholders = ",".join("?" for _ in hashes)
        rows = conn.execute(
            f"SELECT * FROM model_stats WHERE sha256 IN ({placeholders})",
            hashes,
        ).fetchall()
        return {row["sha256"]: dict(row) for row in rows}

    def get_all(self) -> dict[str, dict]:
        """Return all stats keyed by sha256."""
        conn = self._ensure_conn()
        rows = conn.execute("SELECT * FROM model_stats").fetchall()
        return {row["sha256"]: dict(row) for row in rows}

    def count(self) -> int:
        """Return number of rows in model_stats."""
        conn = self._ensure_conn()
        return conn.execute("SELECT COUNT(*) FROM model_stats").fetchone()[0]

    def close(self) -> None:
        if self._conn:
            self._conn.close()
            self._conn = None