codenuk_backend_mine/services/ai-mockup-service/src/migrate_database.py
2025-09-09 11:22:09 +05:30

125 lines
4.6 KiB
Python

#!/usr/bin/env python3
"""
Database migration script to fix TLDraw ID issues
"""
import os
import psycopg2
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
def migrate_database():
"""Migrate the database to fix TLDraw ID issues"""
# Database connection details
db_host = os.getenv('POSTGRES_HOST', 'localhost')
db_user = os.getenv('POSTGRES_USER', 'pipeline_admin')
db_password = os.getenv('POSTGRES_PASSWORD', 'secure_pipeline_2024')
db_port = os.getenv('POSTGRES_PORT', '5433')
db_name = os.getenv('POSTGRES_DB', 'dev_pipeline')
try:
conn = psycopg2.connect(
host=db_host,
user=db_user,
password=db_password,
port=db_port,
database=db_name
)
with conn.cursor() as cur:
print("🔄 Migrating database to fix TLDraw ID issues...")
# Check if tldraw_id column exists
cur.execute("""
SELECT column_name
FROM information_schema.columns
WHERE table_name = 'wireframe_elements'
AND column_name = 'tldraw_id'
""")
if not cur.fetchone():
print(" Adding tldraw_id column...")
cur.execute("ALTER TABLE wireframe_elements ADD COLUMN tldraw_id VARCHAR(255)")
# Check if parent_id is already VARCHAR
cur.execute("""
SELECT data_type
FROM information_schema.columns
WHERE table_name = 'wireframe_elements'
AND column_name = 'parent_id'
""")
column_info = cur.fetchone()
if column_info and column_info[0] == 'uuid':
print(" Converting parent_id from UUID to VARCHAR...")
# Drop the foreign key constraint first
cur.execute("""
ALTER TABLE wireframe_elements
DROP CONSTRAINT IF EXISTS wireframe_elements_parent_id_fkey
""")
# Change the column type
cur.execute("""
ALTER TABLE wireframe_elements
ALTER COLUMN parent_id TYPE VARCHAR(255)
""")
# Update the function
print(" Updating get_wireframe_with_elements function...")
cur.execute("""
CREATE OR REPLACE FUNCTION get_wireframe_with_elements(p_wireframe_id UUID)
RETURNS TABLE(
wireframe_data JSONB,
elements_data JSONB
) AS $$
BEGIN
RETURN QUERY
SELECT
to_jsonb(w.*) as wireframe_data,
COALESCE(
jsonb_agg(
jsonb_build_object(
'id', we.id,
'tldraw_id', we.tldraw_id,
'type', we.element_type,
'data', we.element_data,
'position', we.position,
'size', we.size,
'style', we.style,
'parent_id', we.parent_id,
'z_index', we.z_index
) ORDER BY we.z_index, we.created_at
) FILTER (WHERE we.id IS NOT NULL),
'[]'::jsonb
) as elements_data
FROM wireframes w
LEFT JOIN wireframe_elements we ON w.id = we.wireframe_id
WHERE w.id = p_wireframe_id
GROUP BY w.id, w.user_id, w.project_id, w.name, w.description,
w.device_type, w.dimensions, w.metadata, w.is_active,
w.created_at, w.updated_at;
END;
$$ LANGUAGE plpgsql;
""")
conn.commit()
print("✅ Database migration completed successfully!")
conn.close()
return True
except Exception as e:
print(f"❌ Database migration failed: {e}")
return False
if __name__ == "__main__":
print("🚀 Starting database migration...")
success = migrate_database()
if success:
print("\n✅ Migration completed successfully!")
print("The database now supports TLDraw string IDs properly.")
else:
print("\n❌ Migration failed!")
print("Please check the database connection and try again.")