diff --git a/libs/services/src/ri_cloud_services/sumo_access/grid_access.py b/libs/services/src/ri_cloud_services/sumo_access/grid_access.py index 9c11fbc..fff0c59 100644 --- a/libs/services/src/ri_cloud_services/sumo_access/grid_access.py +++ b/libs/services/src/ri_cloud_services/sumo_access/grid_access.py @@ -18,7 +18,7 @@ ) from ._explorer import get_case_by_uuid -from .grid_types import GridInfo, GridPropertyInfo +from .grid_types import GridDimensions, GridInfo, GridPropertyInfo def get_time_filter(time_or_interval_str: str | None) -> TimeFilter: @@ -62,11 +62,11 @@ def __init__(self, access_token: str, case_uuid: str, ensemble_name: str) -> Non def from_case_uuid(cls, access_token: str, case_uuid: str, ensemble_name: str) -> GridAccess: return cls(access_token=access_token, case_uuid=case_uuid, ensemble_name=ensemble_name) - async def get_available_grid_info_list_async(self) -> list[GridInfo]: - """Return the list of available grids with their realizations.""" + async def get_available_grid_names_async(self) -> list[str]: + """Return the list of available grid names.""" case = get_case_by_uuid(self._access_token, self._case_uuid) - grid_context = case.grids.grids.filter(ensemble=self._ensemble_name) + grid_context = case.grids.filter(ensemble=self._ensemble_name) if await grid_context.length_async() == 0: raise NoDataError( f"No grid tables found for ensemble '{self._ensemble_name}' in case '{self._case_uuid}'", @@ -74,18 +74,42 @@ async def get_available_grid_info_list_async(self) -> list[GridInfo]: ) grid_names = await grid_context.names_async + return grid_names - grid_infos: list[GridInfo] = [] - for grid_name in grid_names: - realization_context = grid_context.filter(name=grid_name, realization=True) - realization_ids = await realization_context.realizationids_async - grid_infos.append( - GridInfo( - name=grid_name, - realizations=sorted(int(r) for r in realization_ids), - ) + async def get_grid_info_async(self, grid_name: str) -> GridInfo: + """Return grid info for the given grid name.""" + case = get_case_by_uuid(self._access_token, self._case_uuid) + + grid_context = case.grids.filter(ensemble=self._ensemble_name, name=grid_name) + if await grid_context.length_async() == 0: + raise NoDataError( + f"No grid tables found for ensemble '{self._ensemble_name}' in case '{self._case_uuid}'", + Service.SUMO, ) - return grid_infos + + per_realization_context = grid_context.filter(realization=True) + + realizations: list[int] = [] + dimensions_per_real: list[GridDimensions] = [] + + # Async for loop + async for grid_real_object in per_realization_context: + if not isinstance(grid_real_object, CPGrid): + raise InvalidDataError(f"Expected CPGrid, got {type(grid_real_object)}", Service.SUMO) + + realizations.append(int(grid_real_object.realization)) + dimensions_per_real.append(self._get_grid_realization_dimensions(grid_real_object)) + + return GridInfo(realizations=realizations, dimensions_per_realization=dimensions_per_real) + + @staticmethod + def _get_grid_realization_dimensions(real_grid_object: CPGrid) -> GridDimensions: + dimensions = GridDimensions( + i_count=real_grid_object.metadata["data"]["spec"]["ncol"], + j_count=real_grid_object.metadata["data"]["spec"]["nrow"], + k_count=real_grid_object.metadata["data"]["spec"]["nlay"], + ) + return dimensions async def get_grid_blob_id_async(self, grid_name: str, realization: int) -> str: """Get the blob ID for the grid data for the given case + ensemble.""" diff --git a/libs/services/src/ri_cloud_services/sumo_access/grid_types.py b/libs/services/src/ri_cloud_services/sumo_access/grid_types.py index c21edfc..9fa193a 100644 --- a/libs/services/src/ri_cloud_services/sumo_access/grid_types.py +++ b/libs/services/src/ri_cloud_services/sumo_access/grid_types.py @@ -5,12 +5,21 @@ from dataclasses import dataclass +@dataclass(frozen=True) +class GridDimensions: + """Dimensions of a grid, including the number of cells in each direction.""" + + i_count: int + j_count: int + k_count: int + + @dataclass(frozen=True) class GridInfo: - """A grid name together with the realizations it is available for.""" + """Realizations of a grid name together with the dimensions of each realization.""" - name: str realizations: list[int] + dimensions_per_realization: list[GridDimensions] @dataclass(frozen=True) diff --git a/ri_cloud_api/primary/routers/grids/converters.py b/ri_cloud_api/primary/routers/grids/converters.py new file mode 100644 index 0000000..0d7dabc --- /dev/null +++ b/ri_cloud_api/primary/routers/grids/converters.py @@ -0,0 +1,25 @@ +from ri_cloud_services.sumo_access.grid_access import GridInfo + +from . import schemas + + +def to_api_grid_info(grid_info: GridInfo) -> list[schemas.GridRealizationInfo]: + """ + Convert a GridInfo to a schemas.GridInfo. + """ + if len(grid_info.realizations) != len(grid_info.dimensions_per_realization): + raise ValueError("Mismatch between realizations and dimensions") + + # Create list of schemas.GridRealizationInfo + realization_infos = [ + schemas.GridRealizationInfo( + realization=realization, + dimensions=schemas.GridDimensions( + iCount=dimensions.i_count, + jCount=dimensions.j_count, + kCount=dimensions.k_count, + ), + ) + for realization, dimensions in zip(grid_info.realizations, grid_info.dimensions_per_realization) + ] + return realization_infos diff --git a/ri_cloud_api/primary/routers/grids/router.py b/ri_cloud_api/primary/routers/grids/router.py index e755692..339b78b 100644 --- a/ri_cloud_api/primary/routers/grids/router.py +++ b/ri_cloud_api/primary/routers/grids/router.py @@ -6,27 +6,44 @@ from __future__ import annotations -from fastapi import APIRouter, Header, Path, Query +from fastapi import APIRouter, Header, HTTPException, Path, Query from ri_cloud_services.sumo_access.grid_access import GridAccess from ri_cloud_api.primary.utils.router_headers import extract_required_token -from . import schemas +from . import converters, schemas router = APIRouter(tags=["grids"]) -@router.get("/cases/{case_uuid}/ensembles/{ensemble_name}/grid_info_list") -async def get_grid_info_list( +@router.get("/cases/{case_uuid}/ensembles/{ensemble_name}/grid_names") +async def get_grid_names( authorization: str | None = Header(None, description="Authorization bearer token for Sumo API"), case_uuid: str = Path(description="Sumo case uuid"), ensemble_name: str = Path(description="Ensemble name"), -) -> list[schemas.GridInfo]: +) -> list[str]: + """List available grid names for the given case + ensemble.""" + access_token = extract_required_token(authorization) + access = GridAccess.from_case_uuid(access_token, case_uuid, ensemble_name) + grid_names = await access.get_available_grid_names_async() + return grid_names + + +@router.get("/cases/{case_uuid}/ensembles/{ensemble_name}/grid_info/{grid_name}") +async def get_grid_info( + authorization: str | None = Header(None, description="Authorization bearer token for Sumo API"), + case_uuid: str = Path(description="Sumo case uuid"), + ensemble_name: str = Path(description="Ensemble name"), + grid_name: str = Path(description="Grid name"), +) -> list[schemas.GridRealizationInfo]: """List available grids, with their realizations, for the given case + ensemble.""" access_token = extract_required_token(authorization) access = GridAccess.from_case_uuid(access_token, case_uuid, ensemble_name) - grids = await access.get_available_grid_info_list_async() - return [schemas.GridInfo(name=g.name, realizations=g.realizations) for g in grids] + grid = await access.get_grid_info_async(grid_name) + try: + return converters.to_api_grid_info(grid) + except ValueError as e: + raise HTTPException(status_code=500, detail=f"Error converting grid info for grid '{grid_name}': {e}") from e @router.get("/cases/{case_uuid}/ensembles/{ensemble_name}/grids/{grid_name}/realizations/{realization}/blob_id") diff --git a/ri_cloud_api/primary/routers/grids/schemas.py b/ri_cloud_api/primary/routers/grids/schemas.py index a29e8e0..b175f0d 100644 --- a/ri_cloud_api/primary/routers/grids/schemas.py +++ b/ri_cloud_api/primary/routers/grids/schemas.py @@ -1,11 +1,17 @@ from pydantic import BaseModel -class GridInfo(BaseModel): - name: str - realizations: list[int] - - class GridPropertyInfo(BaseModel): propertyName: str isoDateOrInterval: str | None = None + + +class GridDimensions(BaseModel): + iCount: int + jCount: int + kCount: int + + +class GridRealizationInfo(BaseModel): + realization: int + dimensions: GridDimensions