from app import app, db

def upgrade_database():
    """Add media columns to Node table"""
    with app.app_context():
        # Check if columns already exist using the new SQLAlchemy 2.0 API
        with db.engine.connect() as conn:
            result = conn.execute(db.text("PRAGMA table_info(node)"))
            columns = result.fetchall()
            column_names = [col[1] for col in columns]
            
            # Add image_url column if it doesn't exist
            if 'image_url' not in column_names:
                conn.execute(db.text("ALTER TABLE node ADD COLUMN image_url VARCHAR(500)"))
                print("Added image_url column to Node table")
            
            # Add youtube_url column if it doesn't exist
            if 'youtube_url' not in column_names:
                conn.execute(db.text("ALTER TABLE node ADD COLUMN youtube_url VARCHAR(500)"))
                print("Added youtube_url column to Node table")
            
            conn.commit()
        
        print("Database migration completed successfully")

if __name__ == "__main__":
    upgrade_database()