389 lines
14 KiB
Python
389 lines
14 KiB
Python
import psycopg2
|
|
from psycopg2 import pool as _pgpool
|
|
import threading
|
|
import json
|
|
|
|
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)",
|
|
]:
|
|
cur.execute(sql)
|
|
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):
|
|
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)
|
|
VALUES (%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);
|
|
""", (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))
|
|
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_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 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
|
|
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
|
|
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."""
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
for idx, fname in enumerate(ordered_filenames):
|
|
cur.execute(
|
|
"UPDATE person SET sort_order = %s WHERE filename = %s",
|
|
(idx, fname)
|
|
)
|
|
conn.commit()
|
|
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)
|