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:
+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
|
||||
Reference in New Issue
Block a user