""" 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 _add_missing_columns(conn: sqlite3.Connection, table: str, columns: dict[str, str]) -> None: """ Additive schema migration: add each column only if the table doesn't already have it. A deployed Space keeps its SQLite file on a persistent volume across redeploys, so the schema has to evolve in place rather than assuming a fresh CREATE TABLE. """ existing = {row["name"] for row in conn.execute(f"PRAGMA table_info({table})").fetchall()} for name, definition in columns.items(): if name not in existing: conn.execute(f"ALTER TABLE {table} ADD COLUMN {name} {definition}") 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' ) """ ) # Columns added after the first release. SQLite has no # "ADD COLUMN IF NOT EXISTS", and an existing Space already has a # populated sprites.sqlite3 on its persistent volume, so each one # is added only when missing rather than recreating the table. _add_missing_columns( conn, "sprites", { # Which segmentation path produced the sprite. Provenance, # not a guess -- see segmentation.DetectedSprite.source. "segmentation_source": "TEXT DEFAULT 'gap'", # 1 while category/style/tags are still a machine # suggestion nobody has reviewed. "auto_tagged": "INTEGER DEFAULT 0", "auto_tag_confidence": "REAL", # train / validation / test, assigned at upload time. "split": "TEXT", }, ) 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 segmentation_source: str = "gap" auto_tagged: bool = False auto_tag_confidence: 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, segmentation_source, auto_tagged, auto_tag_confidence, 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, rec.segmentation_source, 1 if rec.auto_tagged else 0, rec.auto_tag_confidence, 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", "auto_tagged", "split", } 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 select_sprite_ids( scope: str = "sheet", sheet_id: str | None = None, status: str | None = None, sprite_ids: list[str] | None = None, ) -> list[str]: """ Resolve a bulk-operation scope to concrete sprite ids. Scopes: "ids" (an explicit list), "sheet" (everything on one sheet), "status" (everything in one workflow state) or "all". """ if scope == "ids": return list(sprite_ids or []) query = "SELECT id FROM sprites WHERE 1=1" params: list = [] if scope == "sheet": if not sheet_id: return [] query += " AND sheet_id = ?" params.append(sheet_id) if status: query += " AND status = ?" params.append(status) query += " ORDER BY sheet_id, sprite_index" with get_conn() as conn: return [r["id"] for r in conn.execute(query, params).fetchall()] def bulk_update_metadata( sprite_ids: list[str], category: str | None = None, style: str | None = None, add_tags: list[str] | None = None, remove_tags: list[str] | None = None, replace_tags: list[str] | None = None, ) -> int: """ Apply the same category / style / tag change to many sprites in ONE transaction. Tags support add/remove as well as replace, because the realistic bulk edit is additive: a sheet of nothing but tree sprites wants "tree" added to whatever the auto-tagger already worked out per sprite, not every other tag wiped. Replace is there for when the existing tags are genuinely wrong. Anything touched here stops counting as an unreviewed machine suggestion, since a human just decided it. """ if not sprite_ids: return 0 now = time.time() add = [t.strip() for t in (add_tags or []) if t.strip()] remove = {t.strip() for t in (remove_tags or []) if t.strip()} touched = 0 with get_conn() as conn: for sprite_id in sprite_ids: row = conn.execute("SELECT tags FROM sprites WHERE id = ?", (sprite_id,)).fetchone() if row is None: continue sets, values = [], [] if category: sets.append("category = ?") values.append(category) if style: sets.append("style = ?") values.append(style) if replace_tags is not None: tags = [t.strip() for t in replace_tags if t.strip()] elif add or remove: tags = json.loads(row["tags"]) if row["tags"] else [] tags = [t for t in tags if t not in remove] for tag in add: if tag not in tags: tags.append(tag) else: tags = None if tags is not None: sets.append("tags = ?") values.append(json.dumps(tags)) if not sets: continue sets.append("auto_tagged = 0") sets.append("updated_at = ?") values.append(now) values.append(sprite_id) conn.execute(f"UPDATE sprites SET {', '.join(sets)} WHERE id = ?", values) touched += 1 return touched 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 duplicate_clusters(statuses: tuple[str, ...] = ("pending", "accepted")) -> list[list[sqlite3.Row]]: """ Group sprites into clusters of mutual duplicates, following the `is_duplicate_of` chain transitively. Chains matter: a tileset's repeated floor tile is flagged against whichever copy happened to be staged first, so tile #40 may point at #12 which points at #3. Treating those as three unrelated pairs would make "keep one of each" keep three. Only clusters with 2+ members are returned; a sprite with no duplicate at all is not a cluster. """ placeholders = ",".join("?" * len(statuses)) with get_conn() as conn: rows = conn.execute( f"SELECT * FROM sprites WHERE status IN ({placeholders})", statuses ).fetchall() by_id = {r["id"]: r for r in rows} parent: dict[str, str] = {} def find(i: str) -> str: parent.setdefault(i, i) while parent[i] != i: parent[i] = parent.setdefault(parent[i], parent[i]) i = parent[i] return i def union(i: str, j: str) -> None: ri, rj = find(i), find(j) if ri != rj: parent[ri] = rj for r in rows: other = r["is_duplicate_of"] # A sprite may point at one that has since been rejected or # deleted; that link simply drops out of the clustering. if other and other in by_id: union(r["id"], other) grouped: dict[str, list[sqlite3.Row]] = {} for r in rows: if r["id"] in parent: grouped.setdefault(find(r["id"]), []).append(r) return [members for members in grouped.values() if len(members) > 1] 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))