From 63247e85b0bcfbba152fc28f364b76bd6c8e5d71 Mon Sep 17 00:00:00 2001
From: Sam <78538841+spwoodcock@users.noreply.github.com>
Date: Fri, 23 Feb 2024 18:20:57 +0100
Subject: [PATCH] fix: data extract splitting by task refined, update download
 endpoint (#1269)

---
 src/backend/app/db/postgis_utils.py        | 29 +++++++++++-----------
 src/backend/app/projects/project_crud.py   | 16 +++++++++++-
 src/backend/app/projects/project_routes.py |  8 +++++-
 3 files changed, 37 insertions(+), 16 deletions(-)

diff --git a/src/backend/app/db/postgis_utils.py b/src/backend/app/db/postgis_utils.py
index 1b544e6b0c..6523b6959b 100644
--- a/src/backend/app/db/postgis_utils.py
+++ b/src/backend/app/db/postgis_utils.py
@@ -278,34 +278,35 @@ async def split_geojson_by_task_areas(
         WHERE tasks.project_id = :project_id;
 
         -- Retrieve task outlines based on the provided project_id
-        WITH task_outlines AS (
-            SELECT id, outline
-            FROM tasks
-            WHERE project_id = :project_id
-        )
         SELECT
-            task_outlines.id AS task_id,
+            tasks.id AS task_id,
             jsonb_build_object(
                 'type', 'FeatureCollection',
-                'features', jsonb_agg(features.feature)
+                'features', jsonb_agg(feature)
             ) AS task_features
         FROM
-            task_outlines
+            tasks
         LEFT JOIN LATERAL (
-            -- Construct a feature collection with geometries per task area
             SELECT
                 jsonb_build_object(
                     'type', 'Feature',
                     'geometry', ST_AsGeoJSON(temp_features.geometry)::jsonb,
                     'properties', temp_features.properties
                 ) AS feature
-            FROM
-                temp_features
+            FROM (
+                SELECT DISTINCT ON (geometry)
+                    id,
+                    geometry,
+                    properties
+                FROM temp_features
+            ) AS temp_features
             WHERE
-                ST_Within(temp_features.geometry, task_outlines.outline)
-        ) AS features ON true
+                ST_Within(temp_features.geometry, tasks.outline)
+        ) AS feature ON true
+        WHERE
+            tasks.project_id = :project_id
         GROUP BY
-            task_outlines.id;
+            tasks.id;
         """
     )
 
diff --git a/src/backend/app/projects/project_crud.py b/src/backend/app/projects/project_crud.py
index a5db99a6a5..8379657ce6 100644
--- a/src/backend/app/projects/project_crud.py
+++ b/src/backend/app/projects/project_crud.py
@@ -1387,7 +1387,9 @@ async def get_task_geometry(db: Session, project_id: int):
 
 
 async def get_project_features_geojson(
-    db: Session, project: Union[db_models.DbProject, int]
+    db: Session,
+    project: Union[db_models.DbProject, int],
+    task_id: Optional[int] = None,
 ) -> FeatureCollection:
     """Get a geojson of all features for a task."""
     if isinstance(project, int):
@@ -1427,6 +1429,18 @@ async def get_project_features_geojson(
             ),
         )
 
+    # Split by task areas if task_id provided
+    if task_id:
+        split_extract_dict = await split_geojson_by_task_areas(
+            db, data_extract_geojson, project_id
+        )
+        if not split_extract_dict:
+            raise HTTPException(
+                status_code=HTTPStatus.UNPROCESSABLE_ENTITY,
+                detail=(f"Failed to extract geojson for task ({task_id})"),
+            )
+        return split_extract_dict[task_id]
+
     return data_extract_geojson
 
 
diff --git a/src/backend/app/projects/project_routes.py b/src/backend/app/projects/project_routes.py
index 4f29830edf..289cf77d16 100644
--- a/src/backend/app/projects/project_routes.py
+++ b/src/backend/app/projects/project_routes.py
@@ -1017,20 +1017,26 @@ async def download_task_boundaries(
 @router.get("/features/download/")
 async def download_features(
     project_id: int,
+    task_id: Optional[int] = None,
     db: Session = Depends(database.get_db),
     current_user: AuthUser = Depends(mapper),
 ):
     """Downloads the features of a project as a GeoJSON file.
 
+    Can generate a geojson for the entire project, or specific task areas.
+
     Args:
         project_id (int): The id of the project.
+        task_id (int): Specify a specific task area to download for.
         db (Session): The database session, provided automatically.
         current_user (AuthUser): Check if user has MAPPER permission.
 
     Returns:
         Response: The HTTP response object containing the downloaded file.
     """
-    feature_collection = await project_crud.get_project_features_geojson(db, project_id)
+    feature_collection = await project_crud.get_project_features_geojson(
+        db, project_id, task_id
+    )
 
     headers = {
         "Content-Disposition": (