211 lines
7.2 KiB
Python
211 lines
7.2 KiB
Python
import psycopg2
|
|
import json
|
|
|
|
DB_CONFIG = {
|
|
"host": "192.168.1.160",
|
|
"port": 5433,
|
|
"dbname": "dv",
|
|
"user": "dev",
|
|
"password": "dev"
|
|
}
|
|
|
|
def get_db_connection():
|
|
return psycopg2.connect(**DB_CONFIG)
|
|
|
|
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",
|
|
]:
|
|
cur.execute(sql)
|
|
conn.commit()
|
|
finally:
|
|
cur.close()
|
|
conn.close()
|
|
|
|
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):
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
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)
|
|
VALUES (%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);
|
|
""", (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))
|
|
conn.commit()
|
|
finally:
|
|
cur.close()
|
|
conn.close()
|
|
|
|
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()
|
|
conn.close()
|
|
|
|
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
|
|
FROM person WHERE filename = %s
|
|
""", (filename,))
|
|
return cur.fetchone()
|
|
finally:
|
|
cur.close()
|
|
conn.close()
|
|
|
|
def list_persons():
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("""
|
|
SELECT filename, name, group_id, clip_description,
|
|
prompt, pose, sort_order, group_name, hidden, has_background, source_refs
|
|
FROM person
|
|
""")
|
|
return cur.fetchall()
|
|
finally:
|
|
cur.close()
|
|
conn.close()
|
|
|
|
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()
|
|
conn.close()
|
|
|
|
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()
|
|
conn.close()
|
|
|
|
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()
|
|
conn.close()
|
|
|
|
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()
|
|
conn.close()
|
|
|
|
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()
|
|
conn.close()
|
|
|
|
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()
|
|
conn.close()
|
|
|
|
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()
|
|
conn.close()
|
|
|
|
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()
|
|
conn.close()
|