diff --git a/src/aeros_equipment/router.py b/src/aeros_equipment/router.py index 909d501..6a31b00 100644 --- a/src/aeros_equipment/router.py +++ b/src/aeros_equipment/router.py @@ -1,9 +1,9 @@ -from src.aeros_equipment.schema import EquipmentSimulation from src.aeros_equipment.service import get_asset_batch -from typing import List, Optional +from typing import List from fastapi import APIRouter, HTTPException, status +from pydantic import BaseModel from sqlalchemy import select from src.aeros_equipment.model import AerosEquipmentCustomParameterData @@ -11,7 +11,6 @@ from src.auth.service import CurrentUser from src.database.core import DbSession from src.database.service import CommonParameters from src.models import StandardResponse -from src.config import RELIABILITY_SERVICE_API from .schema import CustomParameter, EquipmentPagination from .service import save_default_equipment, get_all @@ -59,14 +58,19 @@ async def get_parameter_oreda( "message": "Success" } -@router.get("/set_equipment_simulation", response_model=StandardResponse[object]) -async def set_equipment_simulation( - db_session: DbSession, location_tags: EquipmentSimulation +class AssetBatchRequest(BaseModel): + repairable_tags: List[str] + non_repairable_tags: List[str] + + +@router.post("/test-asset-batch", response_model=StandardResponse[dict]) +async def test_asset_batch( + db_session: DbSession, payload: AssetBatchRequest ): - equipments = get_asset_batch( - location_tags, - location_tags, - RELIABILITY_SERVICE_API + equipments = await get_asset_batch( + payload.repairable_tags, + payload.non_repairable_tags, + db_session, ) return {"data": equipments, "status": "success", "message": "Success"} \ No newline at end of file diff --git a/src/aeros_equipment/service.py b/src/aeros_equipment/service.py index 6ff5de2..adf1513 100644 --- a/src/aeros_equipment/service.py +++ b/src/aeros_equipment/service.py @@ -225,7 +225,7 @@ async def save_default_equipment(*, db_session: DbSession, project_name: str): ) -def get_asset_batch(location_tags: List[str], nr_location_tags: List[str], +async def get_asset_batch(location_tags: List[str], nr_location_tags: List[str], db_session: DbSession, base_url: str = RELIABILITY_SERVICE_API, timeout: int = 600)-> dict: @@ -255,11 +255,12 @@ def get_asset_batch(location_tags: List[str], nr_location_tags: List[str], WHERE location_tag = ANY(:location_tags) """) - repairable_result = db_session.execute(repairable_query, {"location_tags": location_tags}) + repairable_result = await db_session.execute(repairable_query, {"location_tags": location_tags}) for row in repairable_result: location_tag = row[0] - parameters = json.loads(row[1]) if row[1] else {} + raw_parameters = row[1] + parameters = raw_parameters if isinstance(raw_parameters, dict) else json.loads(raw_parameters) if raw_parameters else {} try: distribution, reldisp1, reldisp2 = get_distribution_from_db(parameters, is_repairable=True) @@ -282,11 +283,12 @@ def get_asset_batch(location_tags: List[str], nr_location_tags: List[str], WHERE location_tag = ANY(:nr_location_tags) """) - non_repairable_result = db_session.execute(non_repairable_query, {"nr_location_tags": nr_location_tags}) + non_repairable_result = await db_session.execute(non_repairable_query, {"nr_location_tags": nr_location_tags}) for row in non_repairable_result: location_tag = row[0] - parameters = json.loads(row[1]) if row[1] else {} + raw_parameters = row[1] + parameters = raw_parameters if isinstance(raw_parameters, dict) else json.loads(raw_parameters) if raw_parameters else {} best_distribution = row[2] try: @@ -439,7 +441,7 @@ async def update_equipment_for_simulation( t.location_tag for t in non_repairable.scalars().all() ] - reliability_data = get_asset_batch( + reliability_data = await get_asset_batch( reg_nodes, non_repairable_location_tags, db_session