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>
This commit is contained in:
+21
-17
@@ -19,17 +19,21 @@ import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ── Import node classes ────────────────────────────────────────────────
|
||||
from .nodes.lora_loader import LoraLoaderRemoteLM, LoraTextLoaderRemoteLM
|
||||
from .nodes.lora_stacker import LoraStackerRemoteLM
|
||||
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
|
||||
# Guard relative imports: only execute when loaded as a proper package
|
||||
# (i.e. inside ComfyUI). Skipped when the file is imported standalone by
|
||||
# pytest or other tools that add the directory directly to sys.path.
|
||||
if __package__:
|
||||
# ── Import node classes ────────────────────────────────────────────────
|
||||
from .nodes.lora_loader import LoraLoaderRemoteLM, LoraTextLoaderRemoteLM
|
||||
from .nodes.lora_stacker import LoraStackerRemoteLM
|
||||
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 = {
|
||||
# ── NODE_CLASS_MAPPINGS (how ComfyUI discovers nodes) ──────────────────
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
LoraLoaderRemoteLM.NAME: LoraLoaderRemoteLM,
|
||||
LoraTextLoaderRemoteLM.NAME: LoraTextLoaderRemoteLM,
|
||||
LoraStackerRemoteLM.NAME: LoraStackerRemoteLM,
|
||||
@@ -39,18 +43,18 @@ NODE_CLASS_MAPPINGS = {
|
||||
SaveImageRemoteLM.NAME: SaveImageRemoteLM,
|
||||
WanVideoLoraSelectRemoteLM.NAME: WanVideoLoraSelectRemoteLM,
|
||||
WanVideoLoraTextSelectRemoteLM.NAME: WanVideoLoraTextSelectRemoteLM,
|
||||
}
|
||||
}
|
||||
|
||||
# ── WEB_DIRECTORY tells ComfyUI where to find our JS extensions ───────
|
||||
WEB_DIRECTORY = "./web/comfyui"
|
||||
# ── WEB_DIRECTORY tells ComfyUI where to find our JS extensions ───────
|
||||
WEB_DIRECTORY = "./web/comfyui"
|
||||
|
||||
# ── Register proxy middleware ──────────────────────────────────────────
|
||||
try:
|
||||
# ── Register proxy middleware ──────────────────────────────────────────
|
||||
try:
|
||||
from server import PromptServer # type: ignore
|
||||
from .proxy import register_proxy
|
||||
|
||||
register_proxy(PromptServer.instance.app)
|
||||
except Exception as exc:
|
||||
except Exception as exc:
|
||||
logger.warning("[LM-Remote] Could not register proxy middleware: %s", exc)
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "WEB_DIRECTORY"]
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "WEB_DIRECTORY"]
|
||||
|
||||
@@ -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"]
|
||||
+135
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user