GTC-3172: Support GADM areas in drivers endpoints

app/crud/geostore.py

@@ -190,22 +190,41 @@ async def get_first_row(sql: Select):
     return row
 
 
-async def get_gadm_geostore(
+async def get_gadm_geostore_id(
     admin_provider: str,
     admin_version: str,
     adm_level: int,
-    simplify: float | None,
     country_id: str,
     region_id: str | None = None,
     subregion_id: str | None = None,
-) -> AdminGeostoreResponse:
-    dv: Tuple[str, str] = await admin_params_to_dataset_version(
-        admin_provider, admin_version
+) -> str:
+    src_table = await get_versioned_dataset(admin_provider, admin_version)
+    columns_etc: List[Column | Label] = [
+        db.column("gfw_geostore_id"),
+    ]
+    row = await _find_first_geostore(
+        adm_level,
+        admin_provider,
+        admin_version,
+        columns_etc,
+        country_id,
+        region_id,
+        src_table,
+        subregion_id,
     )
-    dataset, version = dv
+    return row.gfw_geostore_id
 
-    src_table: Table = db.table(version)
-    src_table.schema = dataset
+
+async def build_gadm_geostore(
+    admin_provider: str,
+    admin_version: str,
+    adm_level: int,
+    simplify: float | None,
+    country_id: str,
+    region_id: str | None = None,
+    subregion_id: str | None = None,
+) -> AdminGeostore:
+    src_table = await get_versioned_dataset(admin_provider, admin_version)
 
     columns_etc: List[Column | Label] = [
         db.column("adm_level"),
@@ -230,16 +249,74 @@ async def get_gadm_geostore(
             )
         )
 
+    row = await _find_first_geostore(
+        adm_level,
+        admin_provider,
+        admin_version,
+        columns_etc,
+        country_id,
+        region_id,
+        src_table,
+        subregion_id,
+    )
+
+    if row.geojson is None:
+        raise GeometryIsNullError(
+            "GeoJSON is None, try reducing or eliminating simplification."
+        )
+
+    return await form_admin_geostore(
+        adm_level=adm_level,
+        admin_version=admin_version,
+        area=float(row.gfw_area__ha),
+        bbox=[float(val) for val in row.gfw_bbox],
+        name=str(row.name),
+        geojson=json.loads(row.geojson),
+        geostore_id=str(row.gfw_geostore_id),
+        level_id=str(row.level_id),
+        simplify=simplify,
+    )
+
+
+async def _find_first_geostore(
+    adm_level,
+    admin_provider,
+    admin_version,
+    columns_etc,
+    country_id,
+    region_id,
+    src_table,
+    subregion_id,
+):
+    sql: Select = db.select(columns_etc).select_from(src_table)
+    sql = await add_where_clauses(
+        adm_level,
+        admin_provider,
+        admin_version,
+        country_id,
+        region_id,
+        sql,
+        subregion_id,
+    )
+    row = await get_first_row(sql)
+    if row is None:
+        raise RecordNotFoundError(
+            f"Admin boundary not found in {admin_provider} version {admin_version}"
+        )
+    return row
+
+
+async def add_where_clauses(
+    adm_level, admin_provider, admin_version, country_id, region_id, sql, subregion_id
+):
     where_clauses: List[TextClause] = [
         db.text("adm_level=:adm_level").bindparams(adm_level=str(adm_level))
     ]
-
     # gid_0 is just a three-character value, but all more specific ids are
     # followed by an underscore (which has to be escaped because normally in
     # SQL an underscore is a wildcard) and a revision number (for which we
     # use an UN-escaped underscore).
     level_id_pattern: str = country_id
-
     if adm_level == 0:  # Special-case to avoid slow LIKE
         where_clauses.append(
             db.text("gid_0=:level_id_pattern").bindparams(
@@ -264,33 +341,38 @@ async def get_gadm_geostore(
                 level_id_pattern=level_id_pattern
             )
         )
-
-    sql: Select = db.select(columns_etc).select_from(src_table)
-
     for clause in where_clauses:
         sql = sql.where(clause)
+    return sql
 
-    row = await get_first_row(sql)
-    if row is None:
-        raise RecordNotFoundError(
-            f"Admin boundary not found in {admin_provider} version {admin_version}"
-        )
 
-    if row.geojson is None:
-        raise GeometryIsNullError(
-            "GeoJSON is None, try reducing or eliminating simplification."
-        )
+async def get_versioned_dataset(admin_provider, admin_version):
+    dv: Tuple[str, str] = await admin_params_to_dataset_version(
+        admin_provider, admin_version
+    )
+    dataset, version = dv
+    src_table: Table = db.table(version)
+    src_table.schema = dataset
+    return src_table
 
-    geostore: AdminGeostore = await form_admin_geostore(
-        adm_level=adm_level,
+
+async def get_gadm_geostore(
+    admin_provider: str,
+    admin_version: str,
+    adm_level: int,
+    simplify: float | None,
+    country_id: str,
+    region_id: str | None = None,
+    subregion_id: str | None = None,
+) -> AdminGeostoreResponse:
+    geostore: AdminGeostore = await build_gadm_geostore(
+        admin_provider=admin_provider,
         admin_version=admin_version,
-        area=float(row.gfw_area__ha),
-        bbox=[float(val) for val in row.gfw_bbox],
-        name=str(row.name),
-        geojson=json.loads(row.geojson),
-        geostore_id=str(row.gfw_geostore_id),
-        level_id=str(row.level_id),
+        adm_level=adm_level,
         simplify=simplify,
+        country_id=country_id,
+        region_id=region_id,
+        subregion_id=subregion_id,
     )
 
     return AdminGeostoreResponse(data=geostore)

app/models/pydantic/datamart.py

@@ -1,14 +1,75 @@
+from abc import ABC, abstractmethod
 from enum import Enum
-from typing import Dict, Optional
+from typing import Dict, Literal, Optional, Union
 from uuid import UUID
 
-from pydantic import Field
+from pydantic import Field, root_validator, validator
 
 from app.models.pydantic.responses import Response
 
+from ...crud.geostore import get_gadm_geostore_id
 from .base import StrictBaseModel
 
 
+class AreaOfInterest(StrictBaseModel, ABC):
+    @abstractmethod
+    async def get_geostore_id(self) -> UUID:
+        """Return the unique identifier for the area of interest."""
+        pass
+
+
+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")
+    provider: str = Field("gadm", title="Administrative Boundary Provider")
+    version: str = Field("4.1", title="Administrative Boundary Version")
+
+    async def get_geostore_id(self) -> UUID:
+        admin_level = (
+            sum(
+                1
+                for field in (self.country, self.region, self.subregion)
+                if field is not None
+            )
+            - 1
+        )
+        geostore_id = await get_gadm_geostore_id(
+            admin_provider=self.provider,
+            admin_version=self.version,
+            adm_level=admin_level,
+            country_id=self.country,
+            region_id=self.region,
+            subregion_id=self.subregion,
+        )
+        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"
     pending = "pending"
@@ -31,7 +92,7 @@ class DataMartResource(StrictBaseModel):
     message: Optional[str] = None
     requested_by: Optional[UUID] = None
     endpoint: str
-    metadata: DataMartMetadata = None
+    metadata: Optional[DataMartMetadata] = None
 
 
 class DataMartResourceLink(StrictBaseModel):
@@ -43,7 +104,9 @@ class DataMartResourceLinkResponse(Response):
 
 
 class TreeCoverLossByDriverIn(StrictBaseModel):
-    geostore_id: UUID
+    aoi: Union[GeostoreAreaOfInterest, AdminAreaOfInterest] = Field(
+        ..., discriminator="type"
+    )
     canopy_cover: int = 30
     dataset_version: Dict[str, str] = {}
 

app/routes/datamart/__init__.py

@@ -0,0 +1,56 @@
+OPENAPI_EXTRA = {
+    "parameters": [
+        {
+            "name": "aoi",
+            "in": "query",
+            "required": True,
+            "style": "deepObject",
+            "explode": True,
+            "examples": {
+                "Geostore Area Of Interest": {
+                    "summary": "Geostore Area Of Interest",
+                    "description": "Custom area",
+                    "value": {
+                        "type": "geostore",
+                        "geostore_id": "637d378f-93a9-4364-bfa8-95b6afd28c3a",
+                    },
+                },
+                "Admin Area Of Interest": {
+                    "summary": "Admin Area Of Interest",
+                    "description": "Administrative Boundary",
+                    "value": {
+                        "type": "admin",
+                        "country": "BRA",
+                        "region": "12",
+                        "subregion": "2",
+                    },
+                },
+            },
+            "description": "The Area of Interest",
+            "schema": {
+                "oneOf": [
+                    {"$ref": "#/components/schemas/GeostoreAreaOfInterest"},
+                    {"$ref": "#/components/schemas/AdminAreaOfInterest"},
+                ]
+            },
+        },
+        {
+            "name": "dataset_version",
+            "in": "query",
+            "required": False,
+            "style": "deepObject",
+            "explode": True,
+            "schema": {
+                "type": "object",
+                "additionalProperties": {"type": "string"},
+            },
+            "example": {
+                "umd_tree_cover_loss": "v1.11",
+                "tsc_tree_cover_loss_drivers": "v2023",
+            },
+            "description": (
+                "Pass dataset version overrides as bracketed query parameters.",
+            ),
+        },
+    ]
+}