🎨 refactor(DataMart): create AreaOfInterest abstraction

First, make the change easy.

Author: gtempus

app/models/pydantic/datamart.py

@@ -1,13 +1,27 @@
 from enum import Enum
-from typing import Dict, Optional
+from typing import Dict, Optional, Union
 from uuid import UUID
+from abc import ABC, abstractmethod
 
 from pydantic import Field
 
 from app.models.pydantic.responses import Response
 
 from .base import StrictBaseModel
 
+class AreaOfInterest(StrictBaseModel, ABC):
+    @abstractmethod
+    def get_geostore_id(self) -> UUID:
+        """Return the unique identifier for the area of interest."""
+        pass
+
+
+class GeostoreAreaOfInterest(AreaOfInterest):
+    geostore_id:UUID = Field(..., title="Geostore ID")
+
+    def get_geostore_id(self) -> UUID:
+        return self.geostore_id
+
 
 class AnalysisStatus(str, Enum):
     saved = "saved"
@@ -43,7 +57,7 @@ class DataMartResourceLinkResponse(Response):
 
 
 class TreeCoverLossByDriverIn(StrictBaseModel):
-    geostore_id: UUID
+    aoi: Union[GeostoreAreaOfInterest]
     canopy_cover: int = 30
     dataset_version: Dict[str, str] = {}
 
@@ -76,3 +90,6 @@ class Config:
 
 class TreeCoverLossByDriverResponse(Response):
     data: TreeCoverLossByDriver
+
+
+

app/routes/datamart/land.py

@@ -28,7 +28,7 @@
     DataMartResourceLinkResponse,
     TreeCoverLossByDriver,
     TreeCoverLossByDriverIn,
-    TreeCoverLossByDriverResponse,
+    TreeCoverLossByDriverResponse, AreaOfInterest, GeostoreAreaOfInterest,
 )
 from app.settings.globals import API_URL
 from app.tasks.datamart.land import (
@@ -61,6 +61,11 @@ def _parse_dataset_versions(request: Request) -> Dict[str, str]:
     # Merge dataset version overrides with default dataset versions
     return DEFAULT_LAND_DATASET_VERSIONS | dataset_versions
 
+
+def _parse_area_of_interest(request: Request) -> AreaOfInterest:
+    aoi_info = request.query_params.getlist('aoi[geostore_id]');
+    return GeostoreAreaOfInterest(geostore_id=aoi_info[0])
+
 @router.get(
     "/tree_cover_loss_by_driver",
     response_class=ORJSONResponse,
@@ -69,6 +74,23 @@ def _parse_dataset_versions(request: Request) -> Dict[str, str]:
     status_code=200,
     openapi_extra={
         "parameters": [
+            {
+                "name": "aoi",
+                "in": "query",
+                "required": True,
+                "style": "deepObject",
+                "explode": True,
+                "example": {
+                    "geostore_id": "637d378f-93a9-4364-bfa8-95b6afd28c3a",
+                },
+                "description": "The Area of Interest",
+                "schema": {
+                    "oneOf": [
+                        {"$ref": "#/components/schemas/GeostoreAreaOfInterest"},
+                        {"$ref": "#/components/schemas/WdpaAreaOfInterest"},
+                    ]
+                }
+            },
             {
                 "name": "dataset_version",
                 "in": "query",
@@ -92,14 +114,14 @@ def _parse_dataset_versions(request: Request) -> Dict[str, str]:
 )
 async def tree_cover_loss_by_driver_search(
     *,
-    geostore_id: UUID = Query(..., title="Geostore ID"),
+    aoi: AreaOfInterest = Depends(_parse_area_of_interest),
     canopy_cover: int = Query(30, alias="canopy_cover", title="Canopy cover percent"),
     dataset_versions: Optional[Dict[str, str]] = Depends(_parse_dataset_versions),
     api_key: APIKey = Depends(get_api_key),
 ):
     """Search if a resource exists for a given geostore and canopy cover."""
     resource_id = _get_resource_id(
-        "tree_cover_loss_by_driver", geostore_id, canopy_cover, dataset_versions
+        "tree_cover_loss_by_driver", aoi.get_geostore_id(), canopy_cover, dataset_versions
     )
 
     # check if it exists
@@ -155,7 +177,7 @@ async def tree_cover_loss_by_driver_post(
 
     # check geostore is valid
     try:
-        await get_geostore(data.geostore_id, GeostoreOrigin.rw)
+        await get_geostore(data.aoi.get_geostore_id(), GeostoreOrigin.rw)
     except HTTPException:
         raise HTTPException(
             status_code=422,
@@ -168,7 +190,7 @@ async def tree_cover_loss_by_driver_post(
     dataset_version = DEFAULT_LAND_DATASET_VERSIONS | data.dataset_version
     resource_id = _get_resource_id(
         "tree_cover_loss_by_driver",
-        data.geostore_id,
+        data.aoi.get_geostore_id(),
         data.canopy_cover,
         dataset_version,
     )
@@ -178,7 +200,7 @@ async def tree_cover_loss_by_driver_post(
     background_tasks.add_task(
         compute_tree_cover_loss_by_driver,
         resource_id,
-        data.geostore_id,
+        data.aoi.get_geostore_id(),
         data.canopy_cover,
         dataset_version,
     )

tests_v2/unit/app/routes/datamart/test_land.py

@@ -32,7 +32,7 @@ async def test_get_tree_cover_loss_by_drivers_not_found(
         origin = payload["domains"][0]
 
         headers = {"origin": origin}
-        params = {"x-api-key": api_key, "geostore_id": geostore, "canopy_cover": 30}
+        params = {"x-api-key": api_key, "aoi[geostore_id]": geostore, "canopy_cover": 30}
 
         response = await async_client.get(
             "/v0/land/tree_cover_loss_by_driver", headers=headers, params=params
@@ -59,7 +59,7 @@ async def test_get_tree_cover_loss_by_drivers_found(
         origin = payload["domains"][0]
 
         headers = {"origin": origin}
-        params = {"x-api-key": api_key, "geostore_id": geostore, "canopy_cover": 30}
+        params = {"x-api-key": api_key, "aoi[geostore_id]": geostore, "canopy_cover": 30}
         resource_id = _get_resource_id(
             "tree_cover_loss_by_driver", geostore, 30, DEFAULT_LAND_DATASET_VERSIONS
         )
@@ -94,7 +94,7 @@ async def test_get_tree_cover_loss_by_drivers_with_overrides(
         origin = payload["domains"][0]
 
         headers = {"origin": origin}
-        params = {"x-api-key": api_key, "geostore_id": geostore, "canopy_cover": 30}
+        params = {"x-api-key": api_key, "aoi[geostore_id]": geostore, "canopy_cover": 30}
         resource_id = _get_resource_id(
             "tree_cover_loss_by_driver",
             geostore,
@@ -143,7 +143,7 @@ async def test_get_tree_cover_loss_by_drivers_with_malformed_overrides(
     params = {"x-api-key": api_key, "geostore_id": geostore, "canopy_cover": 30}
 
     response = await async_client.get(
-        "/v0/land/tree_cover_loss_by_driver?x-api-key={api_key}&geostore_id={geostore_id}&canopy_cover=30&dataset_version[umd_tree_cover_loss]]=v1.8&dataset_version[umd_tree_cover_density_2000]=v1.6",
+        f"/v0/land/tree_cover_loss_by_driver?x-api-key={api_key}&aoi[geostore_id]={geostore}&canopy_cover=30&dataset_version[umd_tree_cover_loss]]=v1.8&dataset_version[umd_tree_cover_density_2000]=v1.6",
         headers=headers,
         params=params,
     )
@@ -166,7 +166,7 @@ async def test_post_tree_cover_loss_by_drivers(
 
     headers = {"origin": origin, "x-api-key": api_key}
     payload = {
-        "geostore_id": geostore,
+        "aoi": {"geostore_id": geostore},
         "canopy_cover": 30,
         "dataset_version": {"umd_tree_cover_loss": "v1.8"},
     }