74 lines
2.1 KiB
Python
74 lines
2.1 KiB
Python
"""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!")
|