""" database.py ----------- Lightweight SQLite persistence layer. Tracks: sheets (uploaded sprite sheets), sprites (every detected/extracted sprite with its full metadata + workflow status), and categories (built-in + user-added custom categories). SQLite is more than sufficient here -- this is a single-Space, low-write- concurrency workload (interactive review by essentially one operator at a time), and it gives us real persistence across app restarts without an external service dependency. """ from __future__ import annotations import json import sqlite3 import threading import time import uuid from contextlib import contextmanager from dataclasses import dataclass, field from config import DB_PATH, DEFAULT_CATEGORIES, DEFAULT_STYLES _lock = threading.Lock() def _connect() -> sqlite3.Connection: conn = sqlite3.connect(str(DB_PATH), check_same_thread=False) conn.row_factory = sqlite3.Row conn.execute("PRAGMA journal_mode=WAL;") return conn @contextmanager def get_conn(): with _lock: conn = _connect() try: yield conn conn.commit() finally: conn.close() def init_db() -> None: with get_conn() as conn: conn.execute( """ CREATE TABLE IF NOT EXISTS sheets ( id TEXT PRIMARY KEY, filename TEXT NOT NULL, width INTEGER, height INTEGER, sprite_count INTEGER DEFAULT 0, created_at REAL, zip_batch_id TEXT ) """ ) conn.execute( """ CREATE TABLE IF NOT EXISTS sprites ( id TEXT PRIMARY KEY, sheet_id TEXT NOT NULL, source_file TEXT NOT NULL, sprite_index INTEGER, x INTEGER, y INTEGER, width INTEGER, height INTEGER, png_path TEXT NOT NULL, png_path_transparent TEXT, sha256 TEXT, phash TEXT, quality_json TEXT, confidence REAL, confidence_reasons TEXT, component_count INTEGER, category TEXT DEFAULT 'unknown', style TEXT DEFAULT 'unknown', tags TEXT DEFAULT '[]', status TEXT DEFAULT 'pending', is_duplicate_of TEXT, duplicate_similarity REAL, created_at REAL, updated_at REAL ) """ ) conn.execute( """ CREATE TABLE IF NOT EXISTS categories ( name TEXT PRIMARY KEY, kind TEXT NOT NULL DEFAULT 'category' ) """ ) conn.execute("CREATE INDEX IF NOT EXISTS idx_sprites_status ON sprites(status)") conn.execute("CREATE INDEX IF NOT EXISTS idx_sprites_sheet ON sprites(sheet_id)") conn.execute("CREATE INDEX IF NOT EXISTS idx_sprites_sha256 ON sprites(sha256)") for cat in DEFAULT_CATEGORIES: conn.execute("INSERT OR IGNORE INTO categories (name, kind) VALUES (?, 'category')", (cat,)) for style in DEFAULT_STYLES: conn.execute("INSERT OR IGNORE INTO categories (name, kind) VALUES (?, 'style')", (style,)) def new_id(prefix: str = "sprite") -> str: return f"{prefix}_{uuid.uuid4().hex[:12]}" # --------------------------------------------------------------------------- # Sheets # --------------------------------------------------------------------------- def insert_sheet(sheet_id: str, filename: str, width: int, height: int, zip_batch_id: str | None = None) -> None: with get_conn() as conn: conn.execute( "INSERT INTO sheets (id, filename, width, height, sprite_count, created_at, zip_batch_id) VALUES (?,?,?,?,?,?,?)", (sheet_id, filename, width, height, 0, time.time(), zip_batch_id), ) def update_sheet_sprite_count(sheet_id: str, count: int) -> None: with get_conn() as conn: conn.execute("UPDATE sheets SET sprite_count = ? WHERE id = ?", (count, sheet_id)) def list_sheets() -> list[sqlite3.Row]: with get_conn() as conn: return conn.execute("SELECT * FROM sheets ORDER BY created_at DESC").fetchall() # --------------------------------------------------------------------------- # Sprites # --------------------------------------------------------------------------- @dataclass class SpriteRecord: id: str sheet_id: str source_file: str sprite_index: int x: int y: int width: int height: int png_path: str png_path_transparent: str | None sha256: str phash: str quality: dict confidence: float confidence_reasons: list[str] component_count: int category: str = "unknown" style: str = "unknown" tags: list[str] = field(default_factory=list) status: str = "pending" is_duplicate_of: str | None = None duplicate_similarity: float | None = None def insert_sprite(rec: SpriteRecord) -> None: now = time.time() with get_conn() as conn: conn.execute( """ INSERT INTO sprites ( id, sheet_id, source_file, sprite_index, x, y, width, height, png_path, png_path_transparent, sha256, phash, quality_json, confidence, confidence_reasons, component_count, category, style, tags, status, is_duplicate_of, duplicate_similarity, created_at, updated_at ) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?) """, ( rec.id, rec.sheet_id, rec.source_file, rec.sprite_index, rec.x, rec.y, rec.width, rec.height, rec.png_path, rec.png_path_transparent, rec.sha256, rec.phash, json.dumps(rec.quality), rec.confidence, json.dumps(rec.confidence_reasons), rec.component_count, rec.category, rec.style, json.dumps(rec.tags), rec.status, rec.is_duplicate_of, rec.duplicate_similarity, now, now, ), ) def update_sprite_status(sprite_id: str, status: str) -> None: with get_conn() as conn: conn.execute("UPDATE sprites SET status = ?, updated_at = ? WHERE id = ?", (status, time.time(), sprite_id)) def update_sprite_fields(sprite_id: str, **fields) -> None: if not fields: return allowed = {"category", "style", "tags", "status", "x", "y", "width", "height", "png_path", "png_path_transparent"} sets = [] values = [] for k, v in fields.items(): if k not in allowed: continue if k == "tags" and isinstance(v, list): v = json.dumps(v) sets.append(f"{k} = ?") values.append(v) if not sets: return sets.append("updated_at = ?") values.append(time.time()) values.append(sprite_id) with get_conn() as conn: conn.execute(f"UPDATE sprites SET {', '.join(sets)} WHERE id = ?", values) def get_sprite(sprite_id: str) -> sqlite3.Row | None: with get_conn() as conn: return conn.execute("SELECT * FROM sprites WHERE id = ?", (sprite_id,)).fetchone() def list_sprites(status: str | None = None, sheet_id: str | None = None, limit: int = 5000) -> list[sqlite3.Row]: query = "SELECT * FROM sprites WHERE 1=1" params: list = [] if status: query += " AND status = ?" params.append(status) if sheet_id: query += " AND sheet_id = ?" params.append(sheet_id) query += " ORDER BY sheet_id, sprite_index LIMIT ?" params.append(limit) with get_conn() as conn: return conn.execute(query, params).fetchall() def existing_hashes_for_dup_check(exclude_status: tuple[str, ...] = ("rejected",)) -> list[dict]: placeholders = ",".join("?" * len(exclude_status)) with get_conn() as conn: rows = conn.execute( f"SELECT id, sha256, phash FROM sprites WHERE status NOT IN ({placeholders})", exclude_status, ).fetchall() return [{"id": r["id"], "sha256": r["sha256"], "phash": r["phash"]} for r in rows] def counts_by_status() -> dict: with get_conn() as conn: rows = conn.execute("SELECT status, COUNT(*) as c FROM sprites GROUP BY status").fetchall() out = {"pending": 0, "accepted": 0, "rejected": 0, "uploaded": 0} for r in rows: out[r["status"]] = r["c"] return out def count_possible_duplicates() -> int: with get_conn() as conn: row = conn.execute( "SELECT COUNT(*) as c FROM sprites WHERE is_duplicate_of IS NOT NULL AND status != 'rejected'" ).fetchone() return row["c"] if row else 0 def mark_duplicate(sprite_id: str, of_sprite_id: str, similarity: float) -> None: with get_conn() as conn: conn.execute( "UPDATE sprites SET is_duplicate_of = ?, duplicate_similarity = ?, updated_at = ? WHERE id = ?", (of_sprite_id, similarity, time.time(), sprite_id), ) def clear_duplicate_flag(sprite_id: str) -> None: with get_conn() as conn: conn.execute( "UPDATE sprites SET is_duplicate_of = NULL, duplicate_similarity = NULL, updated_at = ? WHERE id = ?", (time.time(), sprite_id), ) # --------------------------------------------------------------------------- # Categories # --------------------------------------------------------------------------- def list_categories(kind: str = "category") -> list[str]: with get_conn() as conn: rows = conn.execute("SELECT name FROM categories WHERE kind = ? ORDER BY name", (kind,)).fetchall() return [r["name"] for r in rows] def add_custom_category(name: str, kind: str = "category") -> None: name = name.strip().lower().replace(" ", "_") if not name: return with get_conn() as conn: conn.execute("INSERT OR IGNORE INTO categories (name, kind) VALUES (?, ?)", (name, kind))