WIP

run_migration.py

@@ -0,0 +1,73 @@
+"""Run database migration for training_status fields."""
+import psycopg2
+import os
+
+# Read password from .env file
+password = ""
+try:
+    with open(".env") as f:
+        for line in f:
+            if line.startswith("DB_PASSWORD="):
+                password = line.strip().split("=", 1)[1].strip('"').strip("'")
+                break
+except Exception as e:
+    print(f"Error reading .env: {e}")
+
+print(f"Password found: {bool(password)}")
+
+conn = psycopg2.connect(
+    host="192.168.68.31",
+    port=5432,
+    database="docmaster",
+    user="docmaster",
+    password=password
+)
+conn.autocommit = True
+cur = conn.cursor()
+
+# Add training_status column
+try:
+    cur.execute("ALTER TABLE training_datasets ADD COLUMN training_status VARCHAR(20) DEFAULT NULL")
+    print("Added training_status column")
+except Exception as e:
+    print(f"training_status: {e}")
+
+# Add active_training_task_id column
+try:
+    cur.execute("ALTER TABLE training_datasets ADD COLUMN active_training_task_id UUID DEFAULT NULL")
+    print("Added active_training_task_id column")
+except Exception as e:
+    print(f"active_training_task_id: {e}")
+
+# Create indexes
+try:
+    cur.execute("CREATE INDEX IF NOT EXISTS idx_training_datasets_training_status ON training_datasets(training_status)")
+    print("Created training_status index")
+except Exception as e:
+    print(f"index training_status: {e}")
+
+try:
+    cur.execute("CREATE INDEX IF NOT EXISTS idx_training_datasets_active_training_task_id ON training_datasets(active_training_task_id)")
+    print("Created active_training_task_id index")
+except Exception as e:
+    print(f"index active_training_task_id: {e}")
+
+# Update existing datasets that have been used in completed training tasks to trained status
+try:
+    cur.execute("""
+        UPDATE training_datasets d
+        SET status = 'trained'
+        WHERE d.status = 'ready'
+        AND EXISTS (
+            SELECT 1 FROM training_tasks t
+            WHERE t.dataset_id = d.dataset_id
+            AND t.status = 'completed'
+        )
+    """)
+    print(f"Updated {cur.rowcount} datasets to trained status")
+except Exception as e:
+    print(f"update status: {e}")
+
+cur.close()
+conn.close()
+print("Migration complete!")