feat: add SQLite stats database layer for CivitAI stats

Implements StatsDB class with upsert, batch upsert, get_by_hashes,
get_all, count, and close methods backed by a WAL-mode SQLite database.
Also guards __init__.py relative imports behind __package__ check so
pytest can run without ComfyUI context, and removes empty tests/__init__.py
to prevent pytest Package collector from traversing into the plugin root.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

__init__.py

@@ -19,38 +19,42 @@ import logging
 
 logger = logging.getLogger(__name__)
 
-# ── Import node classes ────────────────────────────────────────────────
+# Guard relative imports: only execute when loaded as a proper package
-from .nodes.lora_loader import LoraLoaderRemoteLM, LoraTextLoaderRemoteLM
+# (i.e. inside ComfyUI). Skipped when the file is imported standalone by
-from .nodes.lora_stacker import LoraStackerRemoteLM
+# pytest or other tools that add the directory directly to sys.path.
-from .nodes.lora_randomizer import LoraRandomizerRemoteLM
+if __package__:
-from .nodes.lora_cycler import LoraCyclerRemoteLM
+    # ── Import node classes ────────────────────────────────────────────────
-from .nodes.lora_pool import LoraPoolRemoteLM
+    from .nodes.lora_loader import LoraLoaderRemoteLM, LoraTextLoaderRemoteLM
-from .nodes.save_image import SaveImageRemoteLM
+    from .nodes.lora_stacker import LoraStackerRemoteLM
-from .nodes.wanvideo import WanVideoLoraSelectRemoteLM, WanVideoLoraTextSelectRemoteLM
+    from .nodes.lora_randomizer import LoraRandomizerRemoteLM
+    from .nodes.lora_cycler import LoraCyclerRemoteLM
+    from .nodes.lora_pool import LoraPoolRemoteLM
+    from .nodes.save_image import SaveImageRemoteLM
+    from .nodes.wanvideo import WanVideoLoraSelectRemoteLM, WanVideoLoraTextSelectRemoteLM
 
-# ── NODE_CLASS_MAPPINGS (how ComfyUI discovers nodes) ──────────────────
+    # ── NODE_CLASS_MAPPINGS (how ComfyUI discovers nodes) ──────────────────
-NODE_CLASS_MAPPINGS = {
+    NODE_CLASS_MAPPINGS = {
-    LoraLoaderRemoteLM.NAME: LoraLoaderRemoteLM,
+        LoraLoaderRemoteLM.NAME: LoraLoaderRemoteLM,
-    LoraTextLoaderRemoteLM.NAME: LoraTextLoaderRemoteLM,
+        LoraTextLoaderRemoteLM.NAME: LoraTextLoaderRemoteLM,
-    LoraStackerRemoteLM.NAME: LoraStackerRemoteLM,
+        LoraStackerRemoteLM.NAME: LoraStackerRemoteLM,
-    LoraRandomizerRemoteLM.NAME: LoraRandomizerRemoteLM,
+        LoraRandomizerRemoteLM.NAME: LoraRandomizerRemoteLM,
-    LoraCyclerRemoteLM.NAME: LoraCyclerRemoteLM,
+        LoraCyclerRemoteLM.NAME: LoraCyclerRemoteLM,
-    LoraPoolRemoteLM.NAME: LoraPoolRemoteLM,
+        LoraPoolRemoteLM.NAME: LoraPoolRemoteLM,
-    SaveImageRemoteLM.NAME: SaveImageRemoteLM,
+        SaveImageRemoteLM.NAME: SaveImageRemoteLM,
-    WanVideoLoraSelectRemoteLM.NAME: WanVideoLoraSelectRemoteLM,
+        WanVideoLoraSelectRemoteLM.NAME: WanVideoLoraSelectRemoteLM,
-    WanVideoLoraTextSelectRemoteLM.NAME: WanVideoLoraTextSelectRemoteLM,
+        WanVideoLoraTextSelectRemoteLM.NAME: WanVideoLoraTextSelectRemoteLM,
-}
+    }
 
-# ── WEB_DIRECTORY tells ComfyUI where to find our JS extensions ───────
+    # ── WEB_DIRECTORY tells ComfyUI where to find our JS extensions ───────
-WEB_DIRECTORY = "./web/comfyui"
+    WEB_DIRECTORY = "./web/comfyui"
 
-# ── Register proxy middleware ──────────────────────────────────────────
+    # ── Register proxy middleware ──────────────────────────────────────────
-try:
+    try:
-    from server import PromptServer  # type: ignore
+        from server import PromptServer  # type: ignore
-    from .proxy import register_proxy
+        from .proxy import register_proxy
 
-    register_proxy(PromptServer.instance.app)
+        register_proxy(PromptServer.instance.app)
-except Exception as exc:
+    except Exception as exc:
-    logger.warning("[LM-Remote] Could not register proxy middleware: %s", exc)
+        logger.warning("[LM-Remote] Could not register proxy middleware: %s", exc)
 
-__all__ = ["NODE_CLASS_MAPPINGS", "WEB_DIRECTORY"]
+    __all__ = ["NODE_CLASS_MAPPINGS", "WEB_DIRECTORY"]

conftest.py

@@ -0,0 +1,3 @@
+# Root-level conftest: prevent pytest from collecting the package __init__.py
+# as a test module (it uses relative imports that require ComfyUI context).
+collect_ignore = ["__init__.py"]

stats_db.py

@@ -0,0 +1,135 @@
+"""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))
+        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."""
+        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

tests/test_stats_db.py

@@ -0,0 +1,99 @@
+"""Tests for the CivitAI stats database layer."""
+import pytest
+from pathlib import Path
+
+
+@pytest.fixture
+def db_path(tmp_path):
+    return tmp_path / "test_stats.db"
+
+
+@pytest.fixture
+def stats_db(db_path):
+    from stats_db import StatsDB
+    db = StatsDB(db_path)
+    db.init()
+    return db
+
+
+def test_init_creates_table(stats_db, db_path):
+    """DB file is created and table exists."""
+    assert db_path.exists()
+
+
+def test_upsert_and_get_single(stats_db):
+    """Insert a stat row and retrieve it by sha256."""
+    stats_db.upsert("abc123", {
+        "civitai_model_id": 1,
+        "civitai_version_id": 10,
+        "download_count": 5000,
+        "rating": 4.8,
+        "rating_count": 120,
+        "thumbs_up_count": 89,
+    })
+    result = stats_db.get_by_hashes(["abc123"])
+    assert "abc123" in result
+    assert result["abc123"]["download_count"] == 5000
+    assert result["abc123"]["rating"] == 4.8
+    assert result["abc123"]["thumbs_up_count"] == 89
+
+
+def test_upsert_updates_existing(stats_db):
+    """Upserting same sha256 updates values."""
+    stats_db.upsert("abc123", {"download_count": 100, "rating": 3.0,
+                                "rating_count": 10, "thumbs_up_count": 5})
+    stats_db.upsert("abc123", {"download_count": 200, "rating": 4.0,
+                                "rating_count": 20, "thumbs_up_count": 15})
+    result = stats_db.get_by_hashes(["abc123"])
+    assert result["abc123"]["download_count"] == 200
+
+
+def test_get_by_hashes_batch(stats_db):
+    """Batch retrieval returns only matching hashes."""
+    for i in range(5):
+        stats_db.upsert(f"hash_{i}", {"download_count": i * 100,
+                                       "rating": 0, "rating_count": 0,
+                                       "thumbs_up_count": 0})
+    result = stats_db.get_by_hashes(["hash_1", "hash_3", "nonexistent"])
+    assert len(result) == 2
+    assert "hash_1" in result
+    assert "hash_3" in result
+    assert "nonexistent" not in result
+
+
+def test_get_by_hashes_empty_input(stats_db):
+    """Empty input returns empty dict."""
+    assert stats_db.get_by_hashes([]) == {}
+
+
+def test_upsert_batch(stats_db):
+    """Batch upsert inserts multiple rows."""
+    rows = [
+        ("h1", {"civitai_model_id": 1, "download_count": 100,
+                "rating": 4.0, "rating_count": 10, "thumbs_up_count": 5}),
+        ("h2", {"civitai_model_id": 2, "download_count": 200,
+                "rating": 3.5, "rating_count": 20, "thumbs_up_count": 15}),
+    ]
+    stats_db.upsert_batch(rows)
+    result = stats_db.get_by_hashes(["h1", "h2"])
+    assert len(result) == 2
+    assert result["h1"]["download_count"] == 100
+    assert result["h2"]["download_count"] == 200
+
+
+def test_get_all_stats(stats_db):
+    """get_all returns every row keyed by sha256."""
+    stats_db.upsert("a", {"download_count": 1, "rating": 0,
+                           "rating_count": 0, "thumbs_up_count": 0})
+    stats_db.upsert("b", {"download_count": 2, "rating": 0,
+                           "rating_count": 0, "thumbs_up_count": 0})
+    result = stats_db.get_all()
+    assert len(result) == 2
+
+
+def test_stats_count(stats_db):
+    """count returns number of rows."""
+    assert stats_db.count() == 0
+    stats_db.upsert("a", {"download_count": 1, "rating": 0,
+                           "rating_count": 0, "thumbs_up_count": 0})
+    assert stats_db.count() == 1