Spaces:
Paused
Paused
| """ | |
| 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 | |
| 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 | |
| # --------------------------------------------------------------------------- | |
| 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)) | |