Untitled

from fastapi import FastAPI, HTTPException
from sqlalchemy.future import select
from sqlalchemy.orm import Session
from src.database.connection.db_session import async_session
from src.models import NodeVariable, NodeNetworkScenario  # Import your models
from typing import List, Dict

app = FastAPI()

@app.get("/custom-scenario/{platform_name}/export/display")
async def export_custom_scenario(platform_name: str) -> Dict[str, List]:
    async with async_session() as session:
        return await get_scenario_data(session, platform_name)

async def get_scenario_data(session: Session, platform_name: str) -> Dict[str, List]:
    stmt = (
        select(
            NodeVariable, 
            NodeNetworkScenario
        )
        .join(NodeNetworkScenario, NodeVariable.scenario_id == NodeNetworkScenario.id)
        .where(NodeNetworkScenario.platform_name == platform_name)
    )
    
    result = await session.execute(stmt)
    data = result.all()

    if not data:
        raise HTTPException(status_code=404, detail="No data found for the given platform name")

    columns = data[0].keys()  # Get column names
    rows = [dict(row) for row in data]  # Convert data into a list of dictionaries

    return {"columns": columns, "rows": rows}