From 2a4511f2b323aa34afe9a1945f3f18bb33a1b9fe Mon Sep 17 00:00:00 2001
From: kshitijrajsharma <skshitizraj@gmail.com>
Date: Sat, 22 Jun 2024 12:58:24 +0545
Subject: [PATCH] refactor(custom-exports-yaml): added geometry and both yaml
 validation

---
 API/custom_exports.py    | 25 +++++++-----------
 src/validation/models.py | 56 ++++++++++++++++++++++------------------
 2 files changed, 41 insertions(+), 40 deletions(-)

diff --git a/API/custom_exports.py b/API/custom_exports.py
index 928a24a7..ed1b9a37 100644
--- a/API/custom_exports.py
+++ b/API/custom_exports.py
@@ -1,3 +1,6 @@
+# Standard library imports
+import json
+
 # Third party imports
 import yaml
 from fastapi import APIRouter, Body, Depends, HTTPException, Request
@@ -9,7 +12,7 @@
 from src.config import DEFAULT_QUEUE_NAME
 from src.config import LIMITER as limiter
 from src.config import RATE_LIMIT_PER_MIN
-from src.validation.models import DynamicCategoriesModel
+from src.validation.models import CategoriesBase, DynamicCategoriesModel
 
 from .api_worker import process_custom_request
 from .auth import AuthUser, UserRole, staff_required
@@ -827,19 +830,7 @@ async def process_custom_requests(
     openapi_extra={
         "requestBody": {
             "content": {
-                "multipart/mixed": {
-                    "schema": {
-                        "type": "object",
-                        "properties": {
-                            "geometry": {"$ref": "#/components/schemas/GeometryModel"},
-                            "yaml_body": {
-                                "type": "string",
-                                "format": "yaml",
-                            },
-                        },
-                        "required": ["geometry", "yaml_body"],
-                    }
-                }
+                "application/x-yaml": {"schema": CategoriesBase.model_json_schema()}
             },
             "required": True,
         },
@@ -849,14 +840,18 @@ async def process_custom_requests(
 @version(1)
 async def process_custom_requests_yaml(
     request: Request,
+    geometry: str,
     user: AuthUser = Depends(staff_required),
 ):
-    print(DynamicCategoriesModel.model_json_schema())
     raw_body = await request.body()
     try:
         data = yaml.safe_load(raw_body)
     except yaml.YAMLError:
         raise HTTPException(status_code=422, detail="Invalid YAML")
+    try:
+        data["geometry"] = json.loads(geometry)
+    except json.decoder.JSONDecodeError as ex:
+        raise HTTPException(status_code=422, detail="Invalid Geometry")
     try:
         validated_data = DynamicCategoriesModel.model_validate(data)
     except ValidationError as e:
diff --git a/src/validation/models.py b/src/validation/models.py
index 69f34212..b1cd9243 100644
--- a/src/validation/models.py
+++ b/src/validation/models.py
@@ -18,17 +18,13 @@
 # <info@hotosm.org>
 """Page contains validation models for application"""
 # Standard library imports
-import json
 from enum import Enum
 from typing import Dict, List, Optional, Union
 
 # Third party imports
-from area import area
 from geojson_pydantic import Feature, FeatureCollection, MultiPolygon, Polygon
-from geojson_pydantic.types import BBox
 from pydantic import BaseModel as PydanticModel
 from pydantic import Field, validator
-from typing_extensions import TypedDict
 
 # Reader imports
 from src.config import (
@@ -560,32 +556,19 @@ def validate_frequency(cls, value):
         return value.strip()
 
 
-class DynamicCategoriesModel(BaseModel, GeometryValidatorMixin):
-    """
-    Model for dynamic categories.
-
-    Fields:
-    - iso3 (Optional[str]): ISO3 Country Code.
-    - dataset (Optional[DatasetConfig]): Dataset Configurations for HDX Upload.
-    - meta (bool): Dumps Meta db in parquet format & HDX config JSON to S3.
-    - hdx_upload (bool): Enable/Disable uploading the dataset to HDX.
-    - categories (List[Dict[str, CategoryModel]]): List of dynamic categories.
-    - geometry (Optional[Union[Polygon, MultiPolygon]]): Custom polygon geometry.
-    """
-
-    iso3: Optional[str] = Field(
-        default=None,
-        description="ISO3 Country Code",
-        min_length=3,
-        max_length=3,
-        example="USA",
-    )
+class CategoriesBase(BaseModel):
     hdx_upload: bool = Field(
         default=False,
         description="Enable/Disable uploading dataset to hdx, False by default",
     )
     dataset: Optional[DatasetConfig] = Field(
-        default=None, description="Dataset Configurations for HDX Upload"
+        default=None,
+        description="Dataset Configurations for HDX Upload",
+        example={
+            "dataset_prefix": "hotosm_project_1",
+            "dataset_folder": "TM",
+            "dataset_title": "Tasking Manger Project 1",
+        },
     )
     queue: Optional[str] = Field(
         default="raw_ondemand",
@@ -613,6 +596,29 @@ class DynamicCategoriesModel(BaseModel, GeometryValidatorMixin):
             }
         ],
     )
+
+
+class DynamicCategoriesModel(CategoriesBase, GeometryValidatorMixin):
+    """
+    Model for dynamic categories.
+
+    Fields:
+    - iso3 (Optional[str]): ISO3 Country Code.
+    - dataset (Optional[DatasetConfig]): Dataset Configurations for HDX Upload.
+    - meta (bool): Dumps Meta db in parquet format & HDX config JSON to S3.
+    - hdx_upload (bool): Enable/Disable uploading the dataset to HDX.
+    - categories (List[Dict[str, CategoryModel]]): List of dynamic categories.
+    - geometry (Optional[Union[Polygon, MultiPolygon]]): Custom polygon geometry.
+    """
+
+    iso3: Optional[str] = Field(
+        default=None,
+        description="ISO3 Country Code",
+        min_length=3,
+        max_length=3,
+        example="USA",
+    )
+
     geometry: Optional[Union[Polygon, MultiPolygon, Feature, FeatureCollection]] = (
         Field(
             default=None,