182 lines
6.5 KiB
Python
182 lines
6.5 KiB
Python
import sqlite3
|
|
|
|
from photoflow.core.models import Asset, Tag
|
|
|
|
SCHEMA_VERSION = 1
|
|
|
|
SCHEMA = """
|
|
CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT);
|
|
|
|
CREATE TABLE IF NOT EXISTS assets (
|
|
immich_id TEXT PRIMARY KEY,
|
|
taken_at TEXT,
|
|
gps_lat REAL, gps_lon REAL,
|
|
place_city TEXT, place_country TEXT,
|
|
type TEXT,
|
|
has_gps INTEGER NOT NULL DEFAULT 0,
|
|
thumb_path TEXT,
|
|
processed INTEGER NOT NULL DEFAULT 0,
|
|
ingested_at TEXT,
|
|
updated_at TEXT
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS tags (
|
|
name TEXT PRIMARY KEY,
|
|
immich_tag_id TEXT,
|
|
count INTEGER NOT NULL DEFAULT 0
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS asset_tags (
|
|
immich_id TEXT NOT NULL,
|
|
tag_name TEXT NOT NULL,
|
|
PRIMARY KEY (immich_id, tag_name)
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS clusters (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
start_at TEXT, end_at TEXT,
|
|
count INTEGER NOT NULL DEFAULT 0,
|
|
suggested_name TEXT,
|
|
confidence REAL NOT NULL DEFAULT 0,
|
|
kind_guess TEXT NOT NULL DEFAULT 'trip',
|
|
status TEXT NOT NULL DEFAULT 'pending',
|
|
decided_name TEXT,
|
|
reviewed_at TEXT,
|
|
notes TEXT
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS cluster_members (
|
|
cluster_id INTEGER NOT NULL,
|
|
immich_id TEXT NOT NULL,
|
|
member_confidence REAL NOT NULL DEFAULT 1.0,
|
|
is_outlier INTEGER NOT NULL DEFAULT 0,
|
|
included INTEGER NOT NULL DEFAULT 1,
|
|
flagged_coverage INTEGER NOT NULL DEFAULT 0,
|
|
PRIMARY KEY (cluster_id, immich_id)
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS writeback_log (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
immich_id TEXT NOT NULL,
|
|
action TEXT NOT NULL,
|
|
tag TEXT,
|
|
result TEXT NOT NULL,
|
|
applied_at TEXT NOT NULL
|
|
);
|
|
|
|
CREATE INDEX IF NOT EXISTS idx_assets_taken_at ON assets(taken_at);
|
|
CREATE INDEX IF NOT EXISTS idx_members_cluster ON cluster_members(cluster_id);
|
|
CREATE INDEX IF NOT EXISTS idx_members_asset ON cluster_members(immich_id);
|
|
"""
|
|
|
|
|
|
class Store:
|
|
def __init__(self, db_path: str):
|
|
self.db_path = db_path
|
|
self._conn = None
|
|
|
|
def connect(self) -> "Store":
|
|
self._conn = sqlite3.connect(self.db_path)
|
|
self._conn.row_factory = sqlite3.Row
|
|
self._conn.execute("PRAGMA foreign_keys=ON")
|
|
self.migrate()
|
|
return self
|
|
|
|
@property
|
|
def conn(self) -> sqlite3.Connection:
|
|
if self._conn is None:
|
|
raise RuntimeError("Store not connected; call connect()")
|
|
return self._conn
|
|
|
|
def migrate(self) -> None:
|
|
self.conn.executescript(SCHEMA)
|
|
if self.get_meta("schema_version") is None:
|
|
self.set_meta("schema_version", str(SCHEMA_VERSION))
|
|
self.conn.commit()
|
|
|
|
def set_meta(self, key: str, value: str) -> None:
|
|
self.conn.execute(
|
|
"INSERT INTO meta(key, value) VALUES(?, ?) "
|
|
"ON CONFLICT(key) DO UPDATE SET value=excluded.value", (key, value))
|
|
self.conn.commit()
|
|
|
|
def get_meta(self, key: str, default=None):
|
|
row = self.conn.execute("SELECT value FROM meta WHERE key=?", (key,)).fetchone()
|
|
return row["value"] if row else default
|
|
|
|
def close(self) -> None:
|
|
if self._conn is not None:
|
|
self._conn.close()
|
|
self._conn = None
|
|
|
|
def upsert_asset(self, a: Asset) -> None:
|
|
has_gps = 1 if (a.gps_lat is not None and a.gps_lon is not None) else 0
|
|
self.conn.execute(
|
|
"""INSERT INTO assets
|
|
(immich_id, taken_at, gps_lat, gps_lon, place_city, place_country,
|
|
type, has_gps, thumb_path, processed, ingested_at, updated_at)
|
|
VALUES (?,?,?,?,?,?,?,?,?,?,?,?)
|
|
ON CONFLICT(immich_id) DO UPDATE SET
|
|
taken_at=excluded.taken_at, gps_lat=excluded.gps_lat,
|
|
gps_lon=excluded.gps_lon, place_city=excluded.place_city,
|
|
place_country=excluded.place_country, type=excluded.type,
|
|
has_gps=excluded.has_gps, thumb_path=excluded.thumb_path,
|
|
ingested_at=excluded.ingested_at, updated_at=excluded.updated_at""",
|
|
(a.immich_id, a.taken_at, a.gps_lat, a.gps_lon, a.place_city,
|
|
a.place_country, a.type, has_gps, a.thumb_path,
|
|
1 if a.processed else 0, a.ingested_at, a.updated_at))
|
|
self.conn.commit()
|
|
|
|
def _asset_from_row(self, r) -> Asset:
|
|
return Asset(
|
|
immich_id=r["immich_id"], taken_at=r["taken_at"],
|
|
gps_lat=r["gps_lat"], gps_lon=r["gps_lon"],
|
|
place_city=r["place_city"], place_country=r["place_country"],
|
|
type=r["type"], has_gps=bool(r["has_gps"]), thumb_path=r["thumb_path"],
|
|
processed=bool(r["processed"]), ingested_at=r["ingested_at"],
|
|
updated_at=r["updated_at"])
|
|
|
|
def get_asset(self, immich_id: str):
|
|
r = self.conn.execute("SELECT * FROM assets WHERE immich_id=?", (immich_id,)).fetchone()
|
|
return self._asset_from_row(r) if r else None
|
|
|
|
def all_assets(self, include_processed: bool = True) -> list:
|
|
sql = "SELECT * FROM assets"
|
|
if not include_processed:
|
|
sql += " WHERE processed=0"
|
|
sql += " ORDER BY taken_at"
|
|
return [self._asset_from_row(r) for r in self.conn.execute(sql)]
|
|
|
|
def assets_in_range(self, start_at: str, end_at: str) -> list:
|
|
return [self._asset_from_row(r) for r in self.conn.execute(
|
|
"SELECT * FROM assets WHERE taken_at>=? AND taken_at<=? ORDER BY taken_at",
|
|
(start_at, end_at))]
|
|
|
|
def mark_processed(self, immich_id: str) -> None:
|
|
self.conn.execute("UPDATE assets SET processed=1 WHERE immich_id=?", (immich_id,))
|
|
self.conn.commit()
|
|
|
|
def set_asset_tags(self, immich_id: str, tag_names: list) -> None:
|
|
self.conn.execute("DELETE FROM asset_tags WHERE immich_id=?", (immich_id,))
|
|
self.conn.executemany(
|
|
"INSERT OR IGNORE INTO asset_tags(immich_id, tag_name) VALUES(?, ?)",
|
|
[(immich_id, n) for n in tag_names])
|
|
self.conn.commit()
|
|
|
|
def asset_tags(self, immich_id: str) -> list:
|
|
return [r["tag_name"] for r in self.conn.execute(
|
|
"SELECT tag_name FROM asset_tags WHERE immich_id=? ORDER BY tag_name",
|
|
(immich_id,))]
|
|
|
|
def upsert_tag(self, name: str, immich_tag_id=None, count: int = 0) -> None:
|
|
self.conn.execute(
|
|
"""INSERT INTO tags(name, immich_tag_id, count) VALUES(?,?,?)
|
|
ON CONFLICT(name) DO UPDATE SET
|
|
immich_tag_id=excluded.immich_tag_id, count=excluded.count""",
|
|
(name, immich_tag_id, count))
|
|
self.conn.commit()
|
|
|
|
def all_tags(self) -> list:
|
|
return [Tag(name=r["name"], immich_tag_id=r["immich_tag_id"], count=r["count"])
|
|
for r in self.conn.execute("SELECT * FROM tags ORDER BY name")]
|