Big-Link-Man/scripts/migrate_add_image_fields.py

102 lines
2.7 KiB
Python

"""
Migration script to add image fields to projects and generated_content tables
Story 7.1: Generate and Insert Images into Articles
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
from src.database.session import db_manager
from sqlalchemy import text
def migrate():
"""Add image fields to projects and generated_content tables"""
session = db_manager.get_session()
try:
print("Starting migration: Add image fields...")
print(" Adding image_theme_prompt to projects table...")
session.execute(text("""
ALTER TABLE projects
ADD COLUMN image_theme_prompt TEXT NULL
"""))
print(" Adding hero_image_url to generated_content table...")
session.execute(text("""
ALTER TABLE generated_content
ADD COLUMN hero_image_url TEXT NULL
"""))
print(" Adding content_images to generated_content table...")
session.execute(text("""
ALTER TABLE generated_content
ADD COLUMN content_images JSON NULL
"""))
session.commit()
print("Migration completed successfully!")
print("\nNew fields added:")
print(" - projects.image_theme_prompt (TEXT, nullable)")
print(" - generated_content.hero_image_url (TEXT, nullable)")
print(" - generated_content.content_images (JSON, nullable)")
except Exception as e:
session.rollback()
print(f"Migration failed: {e}")
raise
finally:
session.close()
def rollback():
"""Rollback migration (remove image fields)"""
session = db_manager.get_session()
try:
print("Rolling back migration: Remove image fields...")
print(" Removing content_images column...")
session.execute(text("""
ALTER TABLE generated_content
DROP COLUMN content_images
"""))
print(" Removing hero_image_url column...")
session.execute(text("""
ALTER TABLE generated_content
DROP COLUMN hero_image_url
"""))
print(" Removing image_theme_prompt column...")
session.execute(text("""
ALTER TABLE projects
DROP COLUMN image_theme_prompt
"""))
session.commit()
print("Rollback completed successfully!")
except Exception as e:
session.rollback()
print(f"Rollback failed: {e}")
raise
finally:
session.close()
if __name__ == "__main__":
if len(sys.argv) > 1 and sys.argv[1] == "rollback":
rollback()
else:
migrate()