import asyncio
from src.database.core import engine
from sqlalchemy import text

async def alter_table():
    async with engine.begin() as conn:
        try:
            await conn.execute(text("ALTER TABLE oh_tr_calculation_data ADD COLUMN plant_results JSONB"))
            print("Added plant_results")
        except Exception as e:
            print(f"plant_results exists: {e}")
            
        try:
            await conn.execute(text("ALTER TABLE oh_tr_calculation_data ADD COLUMN fleet_statistics JSONB"))
            print("Added fleet_statistics")
        except Exception as e:
            print(f"fleet_statistics exists: {e}")
            
        try:
            await conn.execute(text("ALTER TABLE oh_tr_calculation_data ADD COLUMN analysis_metadata JSONB"))
            print("Added analysis_metadata")
        except Exception as e:
            print(f"analysis_metadata exists: {e}")

if __name__ == "__main__":
    asyncio.run(alter_table())