TBNilles
359043ad67
New FastAPI container (port 8189) to download and manage models: - Installed Models, Add/Download (CivitAI/HuggingFace/direct URL), Settings views - Persistent SQLite storage for API keys and download history (./sparkyui-data) - Downloads land in ./models, auto-sorted into ComfyUI's standard subfolders - Default COMFYUI_HOST_PATH and SPARKYUI_DATA_PATH to the project root - Wire docker-compose service, env defaults, gitignore, README docs Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
132 lines
3.9 KiB
Python
132 lines
3.9 KiB
Python
"""SQLite persistence: settings (API keys) and download history/queue."""
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
from typing import Any, Optional
|
|
|
|
from .config import DB_PATH, ensure_dirs
|
|
|
|
# A single connection guarded by a lock keeps things simple and safe across the
|
|
# async event loop + background download tasks (sqlite writes are serialized).
|
|
_lock = threading.Lock()
|
|
_conn: Optional[sqlite3.Connection] = None
|
|
|
|
|
|
def _connect() -> sqlite3.Connection:
|
|
global _conn
|
|
if _conn is None:
|
|
ensure_dirs()
|
|
_conn = sqlite3.connect(str(DB_PATH), check_same_thread=False)
|
|
_conn.row_factory = sqlite3.Row
|
|
_conn.execute("PRAGMA journal_mode=WAL")
|
|
return _conn
|
|
|
|
|
|
def init_db() -> None:
|
|
with _lock:
|
|
conn = _connect()
|
|
conn.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS settings (
|
|
key TEXT PRIMARY KEY,
|
|
value TEXT
|
|
)
|
|
"""
|
|
)
|
|
conn.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS downloads (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
url TEXT NOT NULL,
|
|
source TEXT,
|
|
model_type TEXT,
|
|
filename TEXT,
|
|
dest_path TEXT,
|
|
status TEXT NOT NULL DEFAULT 'queued',
|
|
bytes_done INTEGER NOT NULL DEFAULT 0,
|
|
bytes_total INTEGER NOT NULL DEFAULT 0,
|
|
error TEXT,
|
|
created_at REAL NOT NULL,
|
|
updated_at REAL NOT NULL
|
|
)
|
|
"""
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
# ---- settings -------------------------------------------------------------
|
|
|
|
def get_setting(key: str) -> Optional[str]:
|
|
with _lock:
|
|
row = _connect().execute(
|
|
"SELECT value FROM settings WHERE key = ?", (key,)
|
|
).fetchone()
|
|
return row["value"] if row else None
|
|
|
|
|
|
def set_setting(key: str, value: Optional[str]) -> None:
|
|
with _lock:
|
|
conn = _connect()
|
|
if value is None or value == "":
|
|
conn.execute("DELETE FROM settings WHERE key = ?", (key,))
|
|
else:
|
|
conn.execute(
|
|
"INSERT INTO settings(key, value) VALUES(?, ?) "
|
|
"ON CONFLICT(key) DO UPDATE SET value = excluded.value",
|
|
(key, value),
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
# ---- downloads ------------------------------------------------------------
|
|
|
|
def create_download(url: str, source: str, model_type: str, filename: str,
|
|
dest_path: str) -> int:
|
|
now = time.time()
|
|
with _lock:
|
|
conn = _connect()
|
|
cur = conn.execute(
|
|
"INSERT INTO downloads(url, source, model_type, filename, dest_path, "
|
|
"status, created_at, updated_at) VALUES(?, ?, ?, ?, ?, 'queued', ?, ?)",
|
|
(url, source, model_type, filename, dest_path, now, now),
|
|
)
|
|
conn.commit()
|
|
return int(cur.lastrowid)
|
|
|
|
|
|
def update_download(download_id: int, **fields: Any) -> None:
|
|
if not fields:
|
|
return
|
|
fields["updated_at"] = time.time()
|
|
cols = ", ".join(f"{k} = ?" for k in fields)
|
|
values = list(fields.values()) + [download_id]
|
|
with _lock:
|
|
conn = _connect()
|
|
conn.execute(f"UPDATE downloads SET {cols} WHERE id = ?", values)
|
|
conn.commit()
|
|
|
|
|
|
def get_download(download_id: int) -> Optional[dict]:
|
|
with _lock:
|
|
row = _connect().execute(
|
|
"SELECT * FROM downloads WHERE id = ?", (download_id,)
|
|
).fetchone()
|
|
return dict(row) if row else None
|
|
|
|
|
|
def list_downloads() -> list[dict]:
|
|
with _lock:
|
|
rows = _connect().execute(
|
|
"SELECT * FROM downloads ORDER BY id DESC"
|
|
).fetchall()
|
|
return [dict(r) for r in rows]
|
|
|
|
|
|
def delete_download(download_id: int) -> None:
|
|
with _lock:
|
|
conn = _connect()
|
|
conn.execute("DELETE FROM downloads WHERE id = ?", (download_id,))
|
|
conn.commit()
|