aa

tour-comfy/database.py

@@ -12,21 +12,66 @@ DB_CONFIG = {
 def get_db_connection():
     return psycopg2.connect(**DB_CONFIG)
 
-def upsert_person(filename, filepath=None, name=None, group_id=None, tags=None, embedding=None, clip_description=None):
+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)
-            VALUES (%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);
-        """, (filename, filepath, name, group_id, json.dumps(tags) if tags else None, embedding, clip_description))
+            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()
@@ -36,7 +81,11 @@ def get_person(filename):
     conn = get_db_connection()
     cur = conn.cursor()
     try:
-        cur.execute("SELECT name, group_id, tags, embedding, clip_description, filepath FROM person WHERE filename = %s", (filename,))
+        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()
@@ -46,7 +95,11 @@ def list_persons():
     conn = get_db_connection()
     cur = conn.cursor()
     try:
-        cur.execute("SELECT filename, name, group_id, clip_description FROM person")
+        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()
@@ -56,7 +109,6 @@ def search_similar(embedding, limit=10):
     conn = get_db_connection()
     cur = conn.cursor()
     try:
-        # Convert embedding to string format for pgvector
         embedding_str = "[" + ",".join(map(str, embedding)) + "]"
         cur.execute("""
             SELECT filename, name, group_id, clip_description, embedding <=> %s AS distance
@@ -99,3 +151,60 @@ def get_group_files(group_id):
     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()