feat(AreaOfInterest): add discriminator param for better docs

And improve validation for `AdminAreaOfInterest`.

Author: gtempus

app/models/pydantic/datamart.py

@@ -1,9 +1,9 @@
 from enum import Enum
-from typing import Dict, Optional, Union
+from typing import Dict, Literal, Optional, Union
 from uuid import UUID
 from abc import ABC, abstractmethod
 
-from pydantic import Field
+from pydantic import Field, root_validator, validator
 
 from app.models.pydantic.responses import Response
 
@@ -19,13 +19,15 @@ async def get_geostore_id(self) -> UUID:
 
 
 class GeostoreAreaOfInterest(AreaOfInterest):
+    type: Literal['geostore'] = 'geostore'
     geostore_id:UUID = Field(..., title="Geostore ID")
 
     async def get_geostore_id(self) -> UUID:
         return self.geostore_id
 
 
 class AdminAreaOfInterest(AreaOfInterest):
+    type: Literal['admin'] = 'admin'
     country: str = Field(..., title="ISO Country Code")
     region: Optional[str] = Field(None, title="Region")
     subregion: Optional[str] = Field(None, title="Subregion")
@@ -45,6 +47,23 @@ async def get_geostore_id(self) -> UUID:
         )
         return UUID(geostore_id)
 
+    @root_validator
+    def check_region_subregion(cls, values):
+        region = values.get("region")
+        subregion = values.get("subregion")
+        if subregion is not None and region is None:
+            raise ValueError("region must be specified if subregion is provided")
+        return values
+
+    @validator('provider', pre=True, always=True)
+    def set_provider_default(cls, v):
+        return v or 'gadm'
+
+    @validator('version', pre=True, always=True)
+    def set_version_default(cls, v):
+        return v or '4.1'
+
+
 
 class AnalysisStatus(str, Enum):
     saved = "saved"
@@ -80,7 +99,7 @@ class DataMartResourceLinkResponse(Response):
 
 
 class TreeCoverLossByDriverIn(StrictBaseModel):
-    aoi: Union[GeostoreAreaOfInterest, AdminAreaOfInterest]
+    aoi: Union[GeostoreAreaOfInterest, AdminAreaOfInterest] = Field(..., discriminator='type')
     canopy_cover: int = 30
     dataset_version: Dict[str, str] = {}
 

app/routes/datamart/land.py

@@ -17,6 +17,7 @@
 )
 from fastapi.openapi.models import APIKey
 from fastapi.responses import ORJSONResponse
+from pydantic import ValidationError
 
 from app.crud import datamart as datamart_crud
 from app.errors import RecordNotFoundError
@@ -66,21 +67,26 @@ def _parse_dataset_versions(request: Request) -> Dict[str, str]:
 
 
 def _parse_area_of_interest(request: Request) -> AreaOfInterest:
-    if 'aoi[geostore_id]' in request.query_params:
-        return GeostoreAreaOfInterest(geostore_id=request.query_params['aoi[geostore_id]'])
-
-        # Otherwise, check if the request contains admin area information
-    if 'aoi[country]' in request.query_params:
-        return AdminAreaOfInterest(
-            country=request.query_params['aoi[country]'],
-            region=request.query_params.get('aoi[region]'),
-            subregion=request.query_params.get('aoi[subregion]'),
-            provider=request.query_params.get('aoi[provider]'),
-            version=request.query_params.get('aoi[version]'),
-        )
+    params = request.query_params
+    aoi_type = params.get('aoi[type]')
+    try:
+        if aoi_type == 'geostore':
+            return GeostoreAreaOfInterest(geostore_id=params.get('aoi[geostore_id]', None))
+
+            # Otherwise, check if the request contains admin area information
+        if aoi_type == 'admin':
+            return AdminAreaOfInterest(
+                country=params.get('aoi[country]', None),
+                region=params.get('aoi[region]', None),
+                subregion=params.get('aoi[subregion]', None),
+                provider=params.get('aoi[provider]', None),
+                version=params.get('aoi[version]', None),
+            )
 
-    # If neither type is provided, raise an error
-    raise HTTPException(status_code=422, detail="Invalid Area of Interest parameters")
+        # If neither type is provided, raise an error
+        raise HTTPException(status_code=422, detail="Invalid Area of Interest parameters")
+    except ValidationError as e:
+        raise HTTPException(status_code=422, detail=e.errors())
 
 
 @router.get(
@@ -200,7 +206,7 @@ async def tree_cover_loss_by_driver_post(
     except HTTPException:
         raise HTTPException(
             status_code=422,
-            detail=f"Geostore {data.geostore_id} can't be found or is not valid.",
+            detail=f"Geostore {geostore_id} can't be found or is not valid.",
         )
 
     # create initial Job item as pending

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, "aoi[geostore_id]": geostore, "canopy_cover": 30}
+        params = {"x-api-key": api_key, "aoi[type]": "geostore", "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, "aoi[geostore_id]": geostore, "canopy_cover": 30}
+        params = {"x-api-key": api_key, "aoi[type]": "geostore", "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, "aoi[geostore_id]": geostore, "canopy_cover": 30}
+        params = {"x-api-key": api_key, "aoi[type]": "geostore", "aoi[geostore_id]": geostore, "canopy_cover": 30}
         resource_id = _get_resource_id(
             "tree_cover_loss_by_driver",
             geostore,
@@ -140,10 +140,10 @@ async def test_get_tree_cover_loss_by_drivers_with_malformed_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[type]": "geostore", "aoi[geostore_id]": geostore, "canopy_cover": 30}
 
     response = await async_client.get(
-        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",
+        f"/v0/land/tree_cover_loss_by_driver?dataset_version[umd_tree_cover_loss]]=v1.8&dataset_version[umd_tree_cover_density_2000]=v1.6",
         headers=headers,
         params=params,
     )
@@ -166,7 +166,10 @@ async def test_post_tree_cover_loss_by_drivers(
 
     headers = {"origin": origin, "x-api-key": api_key}
     payload = {
-        "aoi": {"geostore_id": geostore},
+        "aoi": {
+            "type": "geostore",
+            "geostore_id": geostore,
+        },
         "canopy_cover": 30,
         "dataset_version": {"umd_tree_cover_loss": "v1.8"},
     }