From 0f55e0632c015a9b35774a74f13a67632bbaa80d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Anders=20Fredrik=20Ki=C3=A6r?=
 <31612826+anders-kiaer@users.noreply.github.com>
Date: Sun, 1 Oct 2023 17:50:33 +0000
Subject: [PATCH] wip

---
 .../src/backend/primary/routers/explore.py    | 12 ++--
 .../primary/routers/parameters/router.py      |  2 +-
 .../backend/primary/routers/surface/router.py |  4 +-
 .../backend/primary/routers/well/router.py    | 11 ++--
 backend/src/services/sumo_access/_helpers.py  | 57 +++++++++++++++----
 .../services/sumo_access/case_inspector.py    | 41 -------------
 .../src/services/sumo_access/grid_access.py   |  4 +-
 .../sumo_access/inplace_volumetrics_access.py |  4 +-
 .../sumo_access/iteration_inspector.py        | 35 ------------
 .../services/sumo_access/parameter_access.py  |  4 +-
 .../services/sumo_access/seismic_access.py    |  4 +-
 .../services/sumo_access/summary_access.py    |  4 +-
 .../services/sumo_access/surface_access.py    |  4 +-
 .../src/services/sumo_access/table_access.py  |  4 +-
 .../sumo_access/well_completions_access.py    |  4 +-
 15 files changed, 76 insertions(+), 118 deletions(-)
 delete mode 100644 backend/src/services/sumo_access/case_inspector.py
 delete mode 100644 backend/src/services/sumo_access/iteration_inspector.py

diff --git a/backend/src/backend/primary/routers/explore.py b/backend/src/backend/primary/routers/explore.py
index 4de93309a..8ee6bc857 100644
--- a/backend/src/backend/primary/routers/explore.py
+++ b/backend/src/backend/primary/routers/explore.py
@@ -4,10 +4,9 @@
 from pydantic import BaseModel
 
 from src.backend.auth.auth_helper import AuthHelper
-from src.services.sumo_access.case_inspector import CaseInspector
-from src.services.sumo_access.iteration_inspector import IterationInspector
 from src.services.sumo_access.sumo_explore import SumoExplore
 from src.services.utils.authenticated_user import AuthenticatedUser
+from src.services.sumo_access._helpers import SumoEnsemble
 
 router = APIRouter()
 
@@ -89,16 +88,15 @@ def get_ensembles(
 
 
 @router.get("/cases/{case_uuid}/ensembles/{ensemble_name}")
-def get_ensemble_details(
+async def get_ensemble_details(
     authenticated_user: AuthenticatedUser = Depends(AuthHelper.get_authenticated_user),
     case_uuid: str = Path(description="Sumo case uuid"),
     ensemble_name: str = Path(description="Ensemble name"),
 ) -> EnsembleDetails:
     """Get more detailed information for an ensemble"""
-    case_inspector = CaseInspector(authenticated_user.get_sumo_access_token(), case_uuid)
-    case_name = case_inspector.get_case_name()
 
-    iteration_inspector: IterationInspector = case_inspector.create_iteration_inspector(ensemble_name)
-    realizations = iteration_inspector.get_realizations()
+    iteration = await SumoEnsemble.from_case_uuid(authenticated_user.get_sumo_access_token(), case_uuid, ensemble_name)
+    case_name = iteration.get_case_name()
+    realizations = iteration.get_realizations()
 
     return EnsembleDetails(name=ensemble_name, case_name=case_name, case_uuid=case_uuid, realizations=realizations)
diff --git a/backend/src/backend/primary/routers/parameters/router.py b/backend/src/backend/primary/routers/parameters/router.py
index a44bda94d..c117d0623 100644
--- a/backend/src/backend/primary/routers/parameters/router.py
+++ b/backend/src/backend/primary/routers/parameters/router.py
@@ -96,5 +96,5 @@ async def get_sensitivities(
 
     access = await ParameterAccess.from_case_uuid(authenticated_user.get_sumo_access_token(), case_uuid, ensemble_name)
 
-    sensitivities = access.get_parameters_and_sensitivities().sensitivities
+    sensitivities = await access.get_parameters_and_sensitivities().sensitivities
     return sensitivities if sensitivities else []
diff --git a/backend/src/backend/primary/routers/surface/router.py b/backend/src/backend/primary/routers/surface/router.py
index 5c8d1cb1a..4a6aef23f 100644
--- a/backend/src/backend/primary/routers/surface/router.py
+++ b/backend/src/backend/primary/routers/surface/router.py
@@ -5,7 +5,6 @@
 
 
 from src.services.sumo_access.surface_access import SurfaceAccess
-from src.services.sumo_access.case_inspector import CaseInspector
 from src.services.smda_access.stratigraphy_access import StratigraphyAccess
 from src.services.smda_access.stratigraphy_utils import sort_stratigraphic_names_by_hierarchy
 from src.services.smda_access.mocked_drogon_smda_access import _mocked_stratigraphy_access
@@ -13,6 +12,7 @@
 from src.services.utils.authenticated_user import AuthenticatedUser
 from src.services.utils.perf_timer import PerfTimer
 from src.backend.auth.auth_helper import AuthHelper
+from src.services.sumo_access._helpers import SumoCase
 
 
 from . import converters
@@ -38,7 +38,7 @@ async def get_surface_directory(
     )
     sumo_surf_dir = surface_access.get_surface_directory()
 
-    case_inspector = CaseInspector(authenticated_user.get_sumo_access_token(), case_uuid)
+    case_inspector = await SumoCase.from_case_uuid(authenticated_user.get_sumo_access_token(), case_uuid)
     strat_column_identifier = case_inspector.get_stratigraphic_column_identifier()
     strat_access: Union[StratigraphyAccess, _mocked_stratigraphy_access.StratigraphyAccess]
 
diff --git a/backend/src/backend/primary/routers/well/router.py b/backend/src/backend/primary/routers/well/router.py
index a9ede000a..ef9e03579 100644
--- a/backend/src/backend/primary/routers/well/router.py
+++ b/backend/src/backend/primary/routers/well/router.py
@@ -5,10 +5,9 @@
 
 from src.services.smda_access import mocked_drogon_smda_access
 from src.services.smda_access.well_access import WellAccess
-from src.services.sumo_access.case_inspector import CaseInspector
 from src.services.utils.authenticated_user import AuthenticatedUser
 from src.backend.auth.auth_helper import AuthHelper
-
+from src.services.sumo_access._helpers import SumoCase
 from src.services.smda_access.types import WellBoreHeader, WellBoreTrajectory
 
 LOGGER = logging.getLogger(__name__)
@@ -17,7 +16,7 @@
 
 
 @router.get("/well_headers/")
-def get_well_headers(
+async def get_well_headers(
     # fmt:off
     authenticated_user: AuthenticatedUser = Depends(AuthHelper.get_authenticated_user),
     case_uuid: str = Query(description="Sumo case uuid"),
@@ -26,7 +25,7 @@ def get_well_headers(
 ) -> List[WellBoreHeader]:
     """Get well headers for all wells in the field"""
 
-    case_inspector = CaseInspector(authenticated_user.get_sumo_access_token(), case_uuid)
+    case_inspector = await SumoCase.from_case_uuid(authenticated_user.get_sumo_access_token(), case_uuid)
     field_identifier = case_inspector.get_field_identifiers()[0]
     well_access: Union[WellAccess, mocked_drogon_smda_access.WellAccess]
     if field_identifier == "DROGON":
@@ -39,7 +38,7 @@ def get_well_headers(
 
 
 @router.get("/field_well_trajectories/")
-def get_field_well_trajectories(
+async def get_field_well_trajectories(
     # fmt:off
     authenticated_user: AuthenticatedUser = Depends(AuthHelper.get_authenticated_user),
     case_uuid: str = Query(description="Sumo case uuid"), # Should be field identifier?
@@ -47,7 +46,7 @@ def get_field_well_trajectories(
     # fmt:on
 ) -> List[WellBoreTrajectory]:
     """Get well trajectories for field"""
-    case_inspector = CaseInspector(authenticated_user.get_sumo_access_token(), case_uuid)
+    case_inspector = await SumoCase.from_case_uuid(authenticated_user.get_sumo_access_token(), case_uuid)
     field_identifier = case_inspector.get_field_identifiers()[0]
     well_access: Union[WellAccess, mocked_drogon_smda_access.WellAccess]
     if field_identifier == "DROGON":
diff --git a/backend/src/services/sumo_access/_helpers.py b/backend/src/services/sumo_access/_helpers.py
index 7d7084656..837a0d40c 100644
--- a/backend/src/services/sumo_access/_helpers.py
+++ b/backend/src/services/sumo_access/_helpers.py
@@ -1,8 +1,11 @@
 import os
+from typing import List, Sequence, Tuple
 
 from sumo.wrapper import SumoClient
 from fmu.sumo.explorer.objects import CaseCollection, Case
 
+from .queries.case import get_stratigraphic_column_identifier, get_field_identifiers
+
 
 SUMO_ENV = os.getenv("WEBVIZ_SUMO_ENV", "dev")
 
@@ -12,19 +15,53 @@ def create_sumo_client_instance(access_token: str) -> SumoClient:
     return sumo_client
 
 
-class SumoAccess:
-    def __init__(self, sumo_client: SumoClient, case: Case, case_uuid: str, iteration_name: str):
+async def _init_helper(access_token: str, case_uuid: str) -> Tuple[SumoClient, Case]:
+    sumo_client: SumoClient = create_sumo_client_instance(access_token)
+    case_collection = CaseCollection(sumo_client).filter(uuid=case_uuid)
+
+    if await case_collection.length_async() != 1:
+        raise ValueError(f"None or multiple sumo cases found {case_uuid=}")
+
+    case = case_collection[0]
+
+    return sumo_client, case
+
+
+class SumoCase:
+    def __init__(self, sumo_client: SumoClient, case: Case, case_uuid: str):
         self._sumo_client = sumo_client
-        self._case: Case = case
+        self._case = case
         self._case_uuid = case_uuid
-        self._iteration_name: str = iteration_name
 
     @classmethod
-    async def from_case_uuid(cls, access_token: str, case_uuid: str, iteration_name: str):  # type: ignore # wait on Python 3.11
-        sumo_client: SumoClient = create_sumo_client_instance(access_token)
-        case_collection = CaseCollection(sumo_client).filter(uuid=case_uuid)
-        if await case_collection.length_async() != 1:
-            raise ValueError(f"None or multiple sumo cases found {case_uuid=}")
+    async def from_case_uuid(cls, access_token: str, case_uuid: str):  # type: ignore # wait on Python 3.11
+        sumo_client, case = await _init_helper(access_token, case_uuid)
+        return SumoCase(sumo_client=sumo_client, case=case, case_uuid=case_uuid)
+
+    def get_case_name(self) -> str:
+        """Get name of the case"""
+        return self._case.name
 
-        case = case_collection[0]
+    def get_stratigraphic_column_identifier(self) -> str:
+        """Retrieve the stratigraphic column identifier for a case"""
+        return get_stratigraphic_column_identifier(self._sumo_client, self._case_uuid)
+
+    def get_field_identifiers(self) -> List[str]:
+        """Retrieve the field identifiers for a case"""
+        return get_field_identifiers(self._sumo_client, self._case_uuid)
+
+
+class SumoEnsemble(SumoCase):
+    def __init__(self, sumo_client: SumoClient, case: Case, case_uuid: str, iteration_name: str):
+        super().__init__(sumo_client=sumo_client, case=case, case_uuid=case_uuid)
+        self._iteration_name: str = iteration_name
+
+    @classmethod
+    async def from_case_uuid(cls, access_token: str, case_uuid: str, iteration_name: str):  # type: ignore # wait on Python 3.11  # pylint: disable=arguments-differ
+        sumo_client, case = await _init_helper(access_token, case_uuid)
         return cls(sumo_client=sumo_client, case=case, case_uuid=case_uuid, iteration_name=iteration_name)
+
+    def get_realizations(self) -> Sequence[int]:
+        """Get list of realizations for this iteration"""
+        realizations = self._case.get_realizations(self._iteration_name)
+        return sorted([int(real) for real in realizations])
diff --git a/backend/src/services/sumo_access/case_inspector.py b/backend/src/services/sumo_access/case_inspector.py
deleted file mode 100644
index ce4b83d22..000000000
--- a/backend/src/services/sumo_access/case_inspector.py
+++ /dev/null
@@ -1,41 +0,0 @@
-import logging
-from typing import List
-
-from fmu.sumo.explorer.explorer import CaseCollection, Case, SumoClient
-
-from .queries.case import get_stratigraphic_column_identifier, get_field_identifiers
-from ._helpers import create_sumo_client_instance
-from .iteration_inspector import IterationInspector
-
-LOGGER = logging.getLogger(__name__)
-
-
-class CaseInspector:
-    def __init__(self, access_token: str, case_uuid: str):
-        self._sumo_client: SumoClient = create_sumo_client_instance(access_token)
-        self._case_uuid = case_uuid
-
-        case_collection = CaseCollection(self._sumo_client).filter(uuid=self._case_uuid)
-        if len(case_collection) != 1:
-            raise ValueError(f"None or multiple sumo cases found {self._case_uuid=}")
-
-        self._case: Case = case_collection[0]
-
-    def get_case_name(self) -> str:
-        """Get name of the case"""
-        return self._case.name
-
-    def get_stratigraphic_column_identifier(self) -> str:
-        """Retrieve the stratigraphic column identifier for a case"""
-
-        stratigraphic_column_identifier = get_stratigraphic_column_identifier(self._sumo_client, self._case_uuid)
-        return stratigraphic_column_identifier
-
-    def get_field_identifiers(self) -> List[str]:
-        """Retrieve the field identifiers for a case"""
-
-        field_identifiers = get_field_identifiers(self._sumo_client, self._case_uuid)
-        return field_identifiers
-
-    def create_iteration_inspector(self, iteration_name: str) -> IterationInspector:
-        return IterationInspector.from_sumo_case(self._case, iteration_name)
diff --git a/backend/src/services/sumo_access/grid_access.py b/backend/src/services/sumo_access/grid_access.py
index 8e9dffedc..3585122b2 100644
--- a/backend/src/services/sumo_access/grid_access.py
+++ b/backend/src/services/sumo_access/grid_access.py
@@ -6,7 +6,7 @@
 
 from src.services.utils.perf_timer import PerfTimer
 
-from ._helpers import SumoAccess
+from ._helpers import SumoEnsemble
 from .queries.cpgrid import (
     get_grid_geometry_blob_id,
     get_grid_names,
@@ -18,7 +18,7 @@
 LOGGER = logging.getLogger(__name__)
 
 
-class GridAccess(SumoAccess):
+class GridAccess(SumoEnsemble):
     async def grid_model_names(self) -> List[str]:
         """Get a list of grid model names"""
         return await get_grid_names(self._sumo_client, self._case_uuid, self._iteration_name)
diff --git a/backend/src/services/sumo_access/inplace_volumetrics_access.py b/backend/src/services/sumo_access/inplace_volumetrics_access.py
index ba02c8134..ef553e504 100644
--- a/backend/src/services/sumo_access/inplace_volumetrics_access.py
+++ b/backend/src/services/sumo_access/inplace_volumetrics_access.py
@@ -12,7 +12,7 @@
 from fmu.sumo.explorer.objects import TableCollection
 from pydantic import BaseModel
 
-from ._helpers import SumoAccess
+from ._helpers import SumoEnsemble
 from .generic_types import EnsembleScalarResponse
 
 # from fmu.sumo.explorer.objects.table import AggregatedTable
@@ -71,7 +71,7 @@ class Config:
         orm_mode = True
 
 
-class InplaceVolumetricsAccess(SumoAccess):
+class InplaceVolumetricsAccess(SumoEnsemble):
     def get_table_names_and_metadata(self) -> List[InplaceVolumetricsTableMetaData]:
         """Retrieve the available volumetric tables names and corresponding metadata for the case"""
         vol_table_collections: TableCollection = self._case.tables.filter(
diff --git a/backend/src/services/sumo_access/iteration_inspector.py b/backend/src/services/sumo_access/iteration_inspector.py
deleted file mode 100644
index 0b0bfa7c4..000000000
--- a/backend/src/services/sumo_access/iteration_inspector.py
+++ /dev/null
@@ -1,35 +0,0 @@
-from typing import Sequence
-from fmu.sumo.explorer.explorer import CaseCollection, Case, SumoClient
-
-from ._helpers import create_sumo_client_instance
-
-
-class IterationInspector:
-    """
-    Class for inspecting and retrieving iteration (ensemble) information
-    """
-
-    def __init__(self, sumo_case: Case, iteration_name: str):
-        self._case: Case = sumo_case
-        self._iteration_name: str = iteration_name
-
-    @staticmethod
-    def from_case_uuid(access_token: str, case_uuid: str, iteration_name: str) -> "IterationInspector":
-        """Create inspector instance from case UUID"""
-        sumo_client: SumoClient = create_sumo_client_instance(access_token)
-        case_collection = CaseCollection(sumo_client).filter(uuid=case_uuid)
-        if len(case_collection) != 1:
-            raise ValueError(f"None or multiple sumo cases found {case_uuid=}")
-
-        my_sumo_case: Case = case_collection[0]
-        return IterationInspector(sumo_case=my_sumo_case, iteration_name=iteration_name)
-
-    @staticmethod
-    def from_sumo_case(sumo_case: Case, iteration_name: str) -> "IterationInspector":
-        """Create inspector instance from an already discovered Sumo Case object"""
-        return IterationInspector(sumo_case=sumo_case, iteration_name=iteration_name)
-
-    def get_realizations(self) -> Sequence[int]:
-        """Get list of realizations for this iteration"""
-        realizations = self._case.get_realizations(self._iteration_name)
-        return sorted([int(real) for real in realizations])
diff --git a/backend/src/services/sumo_access/parameter_access.py b/backend/src/services/sumo_access/parameter_access.py
index 0c2781936..ea482722a 100644
--- a/backend/src/services/sumo_access/parameter_access.py
+++ b/backend/src/services/sumo_access/parameter_access.py
@@ -5,7 +5,7 @@
 import pandas as pd
 
 from .queries.parameters import get_parameters_for_iteration, SumoEnsembleParameter
-from ._helpers import SumoAccess
+from ._helpers import SumoEnsemble
 from .parameter_types import (
     EnsembleParameter,
     EnsembleParameters,
@@ -17,7 +17,7 @@
 LOGGER = logging.getLogger(__name__)
 
 
-class ParameterAccess(SumoAccess):
+class ParameterAccess(SumoEnsemble):
     def get_parameters_and_sensitivities(self) -> EnsembleParameters:
         """Retrieve parameters for an ensemble"""
 
diff --git a/backend/src/services/sumo_access/seismic_access.py b/backend/src/services/sumo_access/seismic_access.py
index c119b1b12..8053be118 100644
--- a/backend/src/services/sumo_access/seismic_access.py
+++ b/backend/src/services/sumo_access/seismic_access.py
@@ -5,13 +5,13 @@
 from fmu.sumo.explorer import TimeFilter, TimeType
 from fmu.sumo.explorer.objects.cube_collection import CubeCollection
 
-from ._helpers import SumoAccess
+from ._helpers import SumoEnsemble
 from .seismic_types import SeismicCubeMeta, VdsHandle
 
 LOGGER = logging.getLogger(__name__)
 
 
-class SeismicAccess(SumoAccess):
+class SeismicAccess(SumoEnsemble):
     def get_seismic_directory(self) -> List[SeismicCubeMeta]:
         seismic_cube_collection: CubeCollection = self._case.cubes.filter(iteration=self._iteration_name, realization=0)
         seismic_cube_metas: List[SeismicCubeMeta] = []
diff --git a/backend/src/services/sumo_access/summary_access.py b/backend/src/services/sumo_access/summary_access.py
index d3a177154..62261c5db 100644
--- a/backend/src/services/sumo_access/summary_access.py
+++ b/backend/src/services/sumo_access/summary_access.py
@@ -12,7 +12,7 @@
 from src.services.utils.perf_timer import PerfTimer
 
 from ._field_metadata import create_vector_metadata_from_field_meta
-from ._helpers import SumoAccess
+from ._helpers import SumoEnsemble
 from ._resampling import resample_segmented_multi_real_table
 from .generic_types import EnsembleScalarResponse
 from .summary_types import Frequency, VectorInfo, RealizationVector, HistoricalVector, VectorMetadata
@@ -20,7 +20,7 @@
 LOGGER = logging.getLogger(__name__)
 
 
-class SummaryAccess(SumoAccess):
+class SummaryAccess(SumoEnsemble):
     async def get_available_vectors(self) -> List[VectorInfo]:
         timer = PerfTimer()
 
diff --git a/backend/src/services/sumo_access/surface_access.py b/backend/src/services/sumo_access/surface_access.py
index 9ca949b07..db0b188d3 100644
--- a/backend/src/services/sumo_access/surface_access.py
+++ b/backend/src/services/sumo_access/surface_access.py
@@ -9,14 +9,14 @@
 from src.services.utils.perf_timer import PerfTimer
 from src.services.utils.statistic_function import StatisticFunction
 
-from ._helpers import SumoAccess
+from ._helpers import SumoEnsemble
 from .surface_types import SurfaceMeta
 from .generic_types import SumoContent
 
 LOGGER = logging.getLogger(__name__)
 
 
-class SurfaceAccess(SumoAccess):
+class SurfaceAccess(SumoEnsemble):
     def get_surface_directory(self) -> List[SurfaceMeta]:
         surface_collection: SurfaceCollection = self._case.surfaces.filter(
             iteration=self._iteration_name,
diff --git a/backend/src/services/sumo_access/table_access.py b/backend/src/services/sumo_access/table_access.py
index e77908aee..ef19d98ae 100644
--- a/backend/src/services/sumo_access/table_access.py
+++ b/backend/src/services/sumo_access/table_access.py
@@ -3,13 +3,13 @@
 
 import pyarrow as pa
 
-from ._helpers import SumoAccess
+from ._helpers import SumoEnsemble
 from .generic_types import SumoTableSchema
 
 LOGGER = logging.getLogger(__name__)
 
 
-class TableAccess(SumoAccess):
+class TableAccess(SumoEnsemble):
     """Generic access to Sumo tables"""
 
     def get_table_schemas_single_realization(self, realization: int = 0) -> List[SumoTableSchema]:
diff --git a/backend/src/services/sumo_access/well_completions_access.py b/backend/src/services/sumo_access/well_completions_access.py
index 5d67715e3..b759ec475 100644
--- a/backend/src/services/sumo_access/well_completions_access.py
+++ b/backend/src/services/sumo_access/well_completions_access.py
@@ -3,7 +3,7 @@
 
 import pandas as pd
 
-from ._helpers import SumoAccess
+from ._helpers import SumoEnsemble
 
 from .well_completions_types import (
     Completions,
@@ -16,7 +16,7 @@
 )
 
 
-class WellCompletionsAccess(SumoAccess):
+class WellCompletionsAccess(SumoEnsemble):
     """
     Class for accessing and retrieving well completions data
     """