import psycopg2 from psycopg2 import pool as _pgpool import threading import json import time import random DB_CONFIG = { "host": "192.168.1.160", "port": 5433, "dbname": "dv", "user": "dev", "password": "dev" } # A pooled connection is reused across requests instead of paying a fresh TCP + # auth round-trip (~40 ms) on every single query. Under load (a generation # running, a burst of reorders) the old open-per-call design starved the web # threadpool and occasionally tripped Postgres' connection ceiling → 500s. # The pool bounds connections and keeps them warm. If the pool can't be # created or is momentarily exhausted, get_db_connection() falls back to a # direct connect so callers never block or fail. # Cover the web server's sync threadpool (uvicorn/anyio default = 40 workers) # with headroom for background workers, while staying well under Postgres' # max_connections (100). _POOL_MIN = 2 _POOL_MAX = 48 _pool = None _pool_lock = threading.Lock() def _get_pool(): global _pool if _pool is not None: return _pool with _pool_lock: if _pool is None: try: _pool = _pgpool.ThreadedConnectionPool( _POOL_MIN, _POOL_MAX, **DB_CONFIG) except Exception as e: print(f"[db] pool init failed, using direct connections: {e}") _pool = False # sentinel: don't retry on every call return _pool def get_db_connection(): """Return a live DB connection, preferring the pool. Always pair with _put_db_connection() (the existing finally: conn.close() callsites are rewritten to call it) so pooled connections are returned rather than dropped. """ p = _get_pool() if p: try: conn = p.getconn() # Guard against a stale/dead pooled connection. if getattr(conn, "closed", 0): try: p.putconn(conn, close=True) except Exception: pass else: return conn except _pgpool.PoolError: # Pool momentarily exhausted — expected under burst; fall back to a # direct connection silently rather than blocking or failing. pass except Exception as e: print(f"[db] pool getconn failed, direct connect: {e}") return psycopg2.connect(**DB_CONFIG) def _put_db_connection(conn): """Return a connection to the pool (rolling back any open txn) or close it. Safe for both pooled and direct/fallback connections: putconn raises for a connection the pool doesn't own, in which case we just close it. """ if conn is None: return p = _pool try: if p: try: conn.rollback() # clear any aborted/idle-in-txn state before reuse except Exception: pass p.putconn(conn) return except Exception: pass try: conn.close() except Exception: pass def migrate_schema(): """Add new columns to person table if they don't exist. Safe to call repeatedly.""" conn = get_db_connection() cur = conn.cursor() try: for sql in [ "ALTER TABLE person ADD COLUMN IF NOT EXISTS prompt TEXT", "ALTER TABLE person ADD COLUMN IF NOT EXISTS pose TEXT", "ALTER TABLE person ADD COLUMN IF NOT EXISTS sort_order INTEGER", "ALTER TABLE person ADD COLUMN IF NOT EXISTS group_name TEXT", "ALTER TABLE person ADD COLUMN IF NOT EXISTS hidden BOOLEAN DEFAULT FALSE", "ALTER TABLE person ADD COLUMN IF NOT EXISTS has_background BOOLEAN DEFAULT TRUE", "ALTER TABLE person ADD COLUMN IF NOT EXISTS source_refs TEXT", "ALTER TABLE person ADD COLUMN IF NOT EXISTS has_clothing BOOLEAN DEFAULT NULL", "ALTER TABLE person ADD COLUMN IF NOT EXISTS content_type TEXT DEFAULT 'image'", "ALTER TABLE person ADD COLUMN IF NOT EXISTS faceswap_source_video TEXT", "ALTER TABLE person ADD COLUMN IF NOT EXISTS archived BOOLEAN DEFAULT FALSE", "ALTER TABLE person ADD COLUMN IF NOT EXISTS face_embedding vector(512)", "ALTER TABLE person ADD COLUMN IF NOT EXISTS is_source BOOLEAN DEFAULT FALSE", "ALTER TABLE person ADD COLUMN IF NOT EXISTS pose_description TEXT", "ALTER TABLE person ADD COLUMN IF NOT EXISTS pose_skeleton TEXT", "ALTER TABLE person ADD COLUMN IF NOT EXISTS people_count INTEGER DEFAULT NULL", "ALTER TABLE person ADD COLUMN IF NOT EXISTS anatomical_completeness BOOLEAN DEFAULT NULL", "ALTER TABLE person ADD COLUMN IF NOT EXISTS facial_direction TEXT DEFAULT NULL", "ALTER TABLE person ADD COLUMN IF NOT EXISTS objects JSONB DEFAULT NULL", ]: cur.execute(sql) conn.commit() finally: cur.close() _put_db_connection(conn) try: initialize_tags_in_db() except Exception as e: print(f"[db] initialize_tags_in_db error: {e}") def initialize_tags_in_db(): conn = get_db_connection() cur = conn.cursor() try: cur.execute(""" SELECT filename, archived, hidden, is_source, sort_order, has_background, tags FROM person """) rows = cur.fetchall() for filename, archived, hidden, is_source, sort_order, has_background, tags_val in rows: tags_list = [] if tags_val: if isinstance(tags_val, str): try: tags_list = json.loads(tags_val) except Exception: tags_list = [] elif isinstance(tags_val, list): tags_list = tags_val is_archived = bool(archived) if archived is not None else False is_hidden = bool(hidden) if hidden is not None else False is_src = bool(is_source) if is_source is not None else False has_bg = bool(has_background) if has_background is not None else True def add_tag_if_missing(tag): if tag not in tags_list: tags_list.append(tag) if is_archived: add_tag_if_missing("ARCHIVED") if "VISIBLE" in tags_list: tags_list.remove("VISIBLE") else: if not is_hidden: add_tag_if_missing("VISIBLE") if "ARCHIVED" in tags_list: tags_list.remove("ARCHIVED") if is_hidden: add_tag_if_missing("HIDDEN") if "VISIBLE" in tags_list: tags_list.remove("VISIBLE") else: if "HIDDEN" in tags_list: tags_list.remove("HIDDEN") if is_src: add_tag_if_missing("SOURCE") else: if "SOURCE" in tags_list: tags_list.remove("SOURCE") if sort_order == 0: add_tag_if_missing("GROUP_ANCHOR") if "GROUP_MEMBER" in tags_list: tags_list.remove("GROUP_MEMBER") else: add_tag_if_missing("GROUP_MEMBER") if "GROUP_ANCHOR" in tags_list: tags_list.remove("GROUP_ANCHOR") if not has_bg: add_tag_if_missing("BACKGROUND_REMOVED") if "BACKGROUND" in tags_list: tags_list.remove("BACKGROUND") else: add_tag_if_missing("BACKGROUND") if "BACKGROUND_REMOVED" in tags_list: tags_list.remove("BACKGROUND_REMOVED") cur.execute("UPDATE person SET tags = %s WHERE filename = %s", (json.dumps(tags_list), filename)) conn.commit() finally: cur.close() _put_db_connection(conn) def upsert_person(filename, filepath=None, name=None, group_id=None, tags=None, embedding=None, clip_description=None, prompt=None, pose=None, sort_order=None, group_name=None, hidden=None, has_background=None, source_refs=None, has_clothing=None, content_type=None, faceswap_source_video=None, archived=None, face_embedding=None, is_source=None, pose_description=None, pose_skeleton=None, people_count=None, anatomical_completeness=None, facial_direction=None, objects=None): conn = get_db_connection() cur = conn.cursor() face_embedding_str = ("[" + ",".join(map(str, face_embedding)) + "]") if face_embedding is not None else None try: cur.execute(""" INSERT INTO person (filename, filepath, name, group_id, tags, embedding, clip_description, prompt, pose, sort_order, group_name, hidden, has_background, source_refs, has_clothing, content_type, faceswap_source_video, archived, face_embedding, is_source, pose_description, pose_skeleton, people_count, anatomical_completeness, facial_direction, objects) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) ON CONFLICT (filename) DO UPDATE SET filepath = COALESCE(EXCLUDED.filepath, person.filepath), name = COALESCE(EXCLUDED.name, person.name), group_id = COALESCE(EXCLUDED.group_id, person.group_id), tags = COALESCE(EXCLUDED.tags, person.tags), embedding = COALESCE(EXCLUDED.embedding, person.embedding), clip_description = COALESCE(EXCLUDED.clip_description, person.clip_description), prompt = COALESCE(EXCLUDED.prompt, person.prompt), pose = COALESCE(EXCLUDED.pose, person.pose), sort_order = COALESCE(EXCLUDED.sort_order, person.sort_order), group_name = COALESCE(EXCLUDED.group_name, person.group_name), hidden = COALESCE(EXCLUDED.hidden, person.hidden), has_background = COALESCE(EXCLUDED.has_background, person.has_background), source_refs = COALESCE(EXCLUDED.source_refs, person.source_refs), has_clothing = COALESCE(EXCLUDED.has_clothing, person.has_clothing), content_type = COALESCE(EXCLUDED.content_type, person.content_type), faceswap_source_video = COALESCE(EXCLUDED.faceswap_source_video, person.faceswap_source_video), archived = COALESCE(EXCLUDED.archived, person.archived), face_embedding = COALESCE(EXCLUDED.face_embedding, person.face_embedding), is_source = COALESCE(EXCLUDED.is_source, person.is_source), pose_description = COALESCE(EXCLUDED.pose_description, person.pose_description), pose_skeleton = COALESCE(EXCLUDED.pose_skeleton, person.pose_skeleton), people_count = COALESCE(EXCLUDED.people_count, person.people_count), anatomical_completeness = COALESCE(EXCLUDED.anatomical_completeness, person.anatomical_completeness), facial_direction = COALESCE(EXCLUDED.facial_direction, person.facial_direction), objects = COALESCE(EXCLUDED.objects, person.objects); """, (filename, filepath, name, group_id, json.dumps(tags) if tags else None, embedding, clip_description, prompt, pose, sort_order, group_name, hidden, has_background, source_refs, has_clothing, content_type, faceswap_source_video, archived, face_embedding_str, is_source, pose_description, pose_skeleton, people_count, anatomical_completeness, facial_direction, json.dumps(objects) if objects else None)) conn.commit() finally: cur.close() _put_db_connection(conn) def set_archived(filename, archived: bool): conn = get_db_connection() cur = conn.cursor() try: cur.execute("UPDATE person SET archived = %s WHERE filename = %s", (archived, filename)) conn.commit() finally: cur.close() _put_db_connection(conn) def set_filenames_archived(filenames, archived: bool): if not filenames: return [] conn = get_db_connection() cur = conn.cursor() try: cur.execute("UPDATE person SET archived = %s WHERE filename = ANY(%s) RETURNING filename", (archived, list(filenames))) rows = cur.fetchall() updated = [r[0] for r in rows] conn.commit() return updated finally: cur.close() _put_db_connection(conn) def set_hidden(filename, hidden: bool): conn = get_db_connection() cur = conn.cursor() try: cur.execute("UPDATE person SET hidden = %s WHERE filename = %s", (hidden, filename)) conn.commit() finally: cur.close() _put_db_connection(conn) def set_person_tags(filename, tags): conn = get_db_connection() cur = conn.cursor() try: cur.execute("UPDATE person SET tags = %s WHERE filename = %s", (json.dumps(tags) if tags is not None else None, filename)) conn.commit() finally: cur.close() _put_db_connection(conn) def get_person(filename): conn = get_db_connection() cur = conn.cursor() try: cur.execute(""" SELECT name, group_id, tags, embedding, clip_description, filepath, prompt, pose, sort_order, group_name, hidden, has_background, source_refs, has_clothing, is_source, pose_description, pose_skeleton, people_count, anatomical_completeness, facial_direction, objects FROM person WHERE filename = %s """, (filename,)) return cur.fetchone() finally: cur.close() _put_db_connection(conn) def list_persons(include_archived=False): conn = get_db_connection() cur = conn.cursor() try: where = "" if include_archived else "WHERE archived IS NOT TRUE" cur.execute(f""" SELECT filename, name, group_id, clip_description, prompt, pose, sort_order, group_name, hidden, has_background, source_refs, has_clothing, content_type, faceswap_source_video, archived, is_source, tags, pose_description, pose_skeleton, people_count, anatomical_completeness, facial_direction, objects FROM person {where} """) return cur.fetchall() finally: cur.close() _put_db_connection(conn) def search_similar(embedding, limit=10): conn = get_db_connection() cur = conn.cursor() try: embedding_str = "[" + ",".join(map(str, embedding)) + "]" cur.execute(""" SELECT filename, name, group_id, clip_description, embedding <=> %s AS distance FROM person WHERE embedding IS NOT NULL ORDER BY distance ASC LIMIT %s; """, (embedding_str, limit)) return cur.fetchall() finally: cur.close() _put_db_connection(conn) def delete_person(filename): conn = get_db_connection() cur = conn.cursor() try: cur.execute("DELETE FROM person WHERE filename = %s", (filename,)) conn.commit() finally: cur.close() _put_db_connection(conn) def delete_group(group_id): conn = get_db_connection() cur = conn.cursor() try: cur.execute("DELETE FROM person WHERE group_id = %s", (group_id,)) conn.commit() finally: cur.close() _put_db_connection(conn) def get_group_files(group_id): conn = get_db_connection() cur = conn.cursor() try: cur.execute("SELECT filename, filepath FROM person WHERE group_id = %s", (group_id,)) return cur.fetchall() finally: cur.close() _put_db_connection(conn) def set_group_order(group_id, ordered_filenames): """Assign sort_order 0,1,2,... to filenames in the given order, resiliently.""" max_retries = 5 for attempt in range(max_retries): conn = get_db_connection() cur = conn.cursor() try: # Deterministic lock ordering to prevent deadlocks: # Fetch and lock all unique filenames alphabetically unique_filenames = sorted(list(set(ordered_filenames))) if unique_filenames: cur.execute( "SELECT filename FROM person WHERE filename = ANY(%s) ORDER BY filename FOR UPDATE", (unique_filenames,) ) for idx, fname in enumerate(ordered_filenames): cur.execute( "UPDATE person SET sort_order = %s WHERE filename = %s", (idx, fname) ) conn.commit() break except (psycopg2.OperationalError, psycopg2.extensions.TransactionRollbackError) as e: pgcode = getattr(e, 'pgcode', None) if pgcode in ('40P01', '40001') or "deadlock" in str(e).lower(): try: conn.rollback() except Exception: pass if attempt == max_retries - 1: raise sleep_time = (0.1 * (2 ** attempt)) + random.uniform(0.01, 0.1) time.sleep(sleep_time) else: raise finally: cur.close() _put_db_connection(conn) def get_group_order(group_id): """Return [(filename, sort_order), ...] sorted by sort_order NULLS LAST.""" conn = get_db_connection() cur = conn.cursor() try: cur.execute(""" SELECT filename, sort_order FROM person WHERE group_id = %s ORDER BY sort_order NULLS LAST, filename """, (group_id,)) return cur.fetchall() finally: cur.close() _put_db_connection(conn) def set_group_name(group_id, name): """Set group_name for every file in the group.""" conn = get_db_connection() cur = conn.cursor() try: cur.execute("UPDATE person SET group_name = %s WHERE group_id = %s", (name, group_id)) conn.commit() finally: cur.close() _put_db_connection(conn) def get_next_sort_order(group_id): """Return max(sort_order)+1 for the group, or 1 if no sorted members exist.""" conn = get_db_connection() cur = conn.cursor() try: cur.execute( "SELECT COALESCE(MAX(sort_order), 0) + 1 FROM person WHERE group_id = %s", (group_id,) ) return cur.fetchone()[0] finally: cur.close() _put_db_connection(conn) def get_all_group_names(): """Return {group_id: group_name} for groups that have a name set.""" conn = get_db_connection() cur = conn.cursor() try: cur.execute(""" SELECT DISTINCT ON (group_id) group_id, group_name FROM person WHERE group_id IS NOT NULL AND group_name IS NOT NULL """) return {row[0]: row[1] for row in cur.fetchall()} finally: cur.close() _put_db_connection(conn) def get_face_embedding(filename): """Return the face_embedding as a list of floats for a filename, or None.""" conn = get_db_connection() cur = conn.cursor() try: cur.execute("SELECT face_embedding FROM person WHERE filename = %s", (filename,)) row = cur.fetchone() if row and row[0] is not None: val = row[0] # psycopg2 without a pgvector adapter returns vectors as plain strings "[f,f,...]" if isinstance(val, str): return [float(x) for x in val.strip("[]").split(",")] return list(val) return None finally: cur.close() _put_db_connection(conn) def search_similar_face(embedding, limit=12, exclude_group_id=None): """Cosine search on face_embedding (stored only for *_face.png rows). Returns [(filename, group_id, distance), ...] sorted ascending by distance. Rows belonging to exclude_group_id are skipped so a group doesn't match itself. """ conn = get_db_connection() cur = conn.cursor() try: embedding_str = "[" + ",".join(map(str, embedding)) + "]" if exclude_group_id: cur.execute(""" SELECT filename, group_id, face_embedding <=> %s AS distance FROM person WHERE face_embedding IS NOT NULL AND (group_id IS NULL OR group_id != %s) ORDER BY distance ASC LIMIT %s """, (embedding_str, exclude_group_id, limit)) else: cur.execute(""" SELECT filename, group_id, face_embedding <=> %s AS distance FROM person WHERE face_embedding IS NOT NULL ORDER BY distance ASC LIMIT %s """, (embedding_str, limit)) return cur.fetchall() finally: cur.close() _put_db_connection(conn)