Refactoring - Added utility to simplify drogon identifier comparisons

backend_py/primary/primary/routers/polygons/router.py

@@ -11,6 +11,7 @@
 from primary.services.sumo_access.case_inspector import CaseInspector
 from primary.services.sumo_access.polygons_access import PolygonsAccess
 from primary.services.utils.authenticated_user import AuthenticatedUser
+from primary.utils.drogon import is_drogon_identifier
 
 from . import converters, schemas
 
@@ -38,7 +39,7 @@ async def get_polygons_directory(
     strat_column_identifier = await case_inspector.get_stratigraphic_column_identifier_async()
     smda_access: Union[SmdaAccess, DrogonSmdaAccess]
 
-    if strat_column_identifier == "DROGON_HAS_NO_STRATCOLUMN":
+    if is_drogon_identifier(strat_column_identifier=strat_column_identifier):
         smda_access = DrogonSmdaAccess()
     else:
         smda_access = SmdaAccess(authenticated_user.get_smda_access_token(), field_identifier=field_identifiers[0])

backend_py/primary/primary/routers/surface/router.py

@@ -17,6 +17,7 @@
 from primary.services.surface_query_service.surface_query_service import batch_sample_surface_in_points_async
 from primary.services.surface_query_service.surface_query_service import RealizationSampleResult
 from primary.utils.response_perf_metrics import ResponsePerfMetrics
+from primary.utils.drogon import is_drogon_identifier
 
 from . import converters
 from . import schemas
@@ -324,7 +325,7 @@ async def _get_stratigraphic_units_for_case_async(
     perf_metrics.record_lap("get-strat-ident")
 
     smda_access: SmdaAccess | DrogonSmdaAccess
-    if strat_column_identifier == "DROGON_HAS_NO_STRATCOLUMN":
+    if is_drogon_identifier(strat_column_identifier=strat_column_identifier):
         smda_access = DrogonSmdaAccess()
     else:
         smda_access = SmdaAccess(authenticated_user.get_smda_access_token(), field_identifier=field_identifiers[0])

backend_py/primary/primary/routers/well/router.py

@@ -7,6 +7,7 @@
 from primary.services.smda_access import SmdaAccess
 from primary.services.utils.authenticated_user import AuthenticatedUser
 from primary.auth.auth_helper import AuthHelper
+from primary.utils.drogon import is_drogon_identifier
 
 from primary.services.ssdl_access.well_access import WellAccess as SsdlWellAccess
 
@@ -27,7 +28,7 @@ async def get_drilled_wellbore_headers(
 ) -> List[schemas.WellboreHeader]:
     """Get wellbore headers for all wells in the field"""
     well_access: Union[SmdaAccess, DrogonSmdaAccess]
-    if field_identifier == "DROGON":
+    if is_drogon_identifier(field_identifier=field_identifier):
         # Handle DROGON
         well_access = DrogonSmdaAccess()
     else:
@@ -48,7 +49,7 @@ async def get_well_trajectories(
 ) -> List[schemas.WellboreTrajectory]:
     """Get well trajectories for field"""
     well_access: Union[SmdaAccess, DrogonSmdaAccess]
-    if field_identifier == "DROGON":
+    if is_drogon_identifier(field_identifier=field_identifier):
         # Handle DROGON
         well_access = DrogonSmdaAccess()
     else:
@@ -72,7 +73,8 @@ async def get_wellbore_pick_identifiers(
 ) -> List[str]:
     """Get wellbore pick identifiers for field and stratigraphic column"""
     well_access: Union[SmdaAccess, DrogonSmdaAccess]
-    if field_identifier == "DROGON":
+
+    if is_drogon_identifier(strat_column_identifier=strat_column_identifier):
         # Handle DROGON
         well_access = DrogonSmdaAccess()
 
@@ -95,7 +97,7 @@ async def get_wellbore_picks_for_pick_identifier(
 ) -> List[schemas.WellborePick]:
     """Get wellbore picks for field and pick identifier"""
     well_access: Union[SmdaAccess, DrogonSmdaAccess]
-    if field_identifier == "DROGON":
+    if is_drogon_identifier(field_identifier=field_identifier):
         # Handle DROGON
         well_access = DrogonSmdaAccess()
 
@@ -116,7 +118,7 @@ async def get_wellbore_picks_for_wellbore(
 ) -> List[schemas.WellborePick]:
     """Get wellbore picks for field and pick identifier"""
     well_access: Union[SmdaAccess, DrogonSmdaAccess]
-    if field_identifier == "DROGON":
+    if is_drogon_identifier(wellbore_uuid=wellbore_uuid):
         # Handle DROGON
         well_access = DrogonSmdaAccess()
 
@@ -137,7 +139,7 @@ async def get_wellbore_completions(
     """Get well bore completions for a single well bore"""
 
     # Handle DROGON
-    if wellbore_uuid in ["drogon_horizontal", "drogon_vertical"]:
+    if is_drogon_identifier(wellbore_uuid=wellbore_uuid):
         return []
 
     well_access = SsdlWellAccess(authenticated_user.get_ssdl_access_token())
@@ -159,7 +161,7 @@ async def get_wellbore_casings(
     """Get well bore casings for a single well bore"""
 
     # Handle DROGON
-    if wellbore_uuid in ["drogon_horizontal", "drogon_vertical"]:
+    if is_drogon_identifier(wellbore_uuid=wellbore_uuid):
         return []
 
     well_access = SsdlWellAccess(authenticated_user.get_ssdl_access_token())
@@ -179,7 +181,7 @@ async def get_wellbore_perforations(
     """Get well bore casing for a single well bore"""
 
     # Handle DROGON
-    if wellbore_uuid in ["drogon_horizontal", "drogon_vertical"]:
+    if is_drogon_identifier(wellbore_uuid=wellbore_uuid):
         return []
 
     well_access = SsdlWellAccess(authenticated_user.get_ssdl_access_token())
@@ -202,7 +204,7 @@ async def get_wellbore_log_curve_headers(
     """Get all log curve headers for a single well bore"""
 
     # Handle DROGON
-    if wellbore_uuid in ["drogon_horizontal", "drogon_vertical"]:
+    if is_drogon_identifier(wellbore_uuid=wellbore_uuid):
         return []
 
     well_access = SsdlWellAccess(authenticated_user.get_ssdl_access_token())
@@ -228,7 +230,7 @@ async def get_log_curve_data(
     """Get log curve data"""
 
     # Handle DROGON
-    if wellbore_uuid in ["drogon_horizontal", "drogon_vertical"]:
+    if is_drogon_identifier(wellbore_uuid=wellbore_uuid):
         raise NotImplementedError("DROGON log curve data not implemented")
 
     well_access = SsdlWellAccess(authenticated_user.get_ssdl_access_token())

backend_py/primary/primary/utils/drogon.py

@@ -0,0 +1,19 @@
+from typing import Optional
+
+
+def is_drogon_identifier(
+    field_identifier: Optional[str] = None,
+    wellbore_uuid: Optional[str] = None,
+    strat_column_identifier: Optional[str] = None,
+) -> bool:
+    """
+    Checks if an element's identifier is for the drogon mock.
+    """
+    if field_identifier == "DROGON":
+        return True
+    if wellbore_uuid in ["drogon_horizontal", "drogon_vertical"]:
+        return True
+    if strat_column_identifier == "DROGON_HAS_NO_STRATCOLUMN":
+        return True
+
+    return False