From fd242416a24b5fc87e2607ab9f336bfab507d2e6 Mon Sep 17 00:00:00 2001
From: kshitijrajsharma <skshitizraj@gmail.com>
Date: Tue, 6 Aug 2024 08:02:34 +0545
Subject: [PATCH] feat(streaming-response-for-api): adds streaming response for
 current snapshot plain

---
 API/raw_data.py | 48 +++++++++++++++++++++++++++++++++++++++---------
 src/app.py      |  2 +-
 2 files changed, 40 insertions(+), 10 deletions(-)

diff --git a/API/raw_data.py b/API/raw_data.py
index 582c81aa..eaa8f3cf 100644
--- a/API/raw_data.py
+++ b/API/raw_data.py
@@ -21,12 +21,14 @@
 """
 # Standard library imports
 import json
+from typing import AsyncGenerator
 
 # Third party imports
+import orjson
 import redis
 from area import area
 from fastapi import APIRouter, Body, Depends, HTTPException, Request
-from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 from fastapi_versioning import version
 
 # Reader imports
@@ -39,6 +41,7 @@
 )
 from src.config import LIMITER as limiter
 from src.config import RATE_LIMIT_PER_MIN as export_rate_limit
+from src.query_builder.builder import raw_currentdata_extraction_query
 from src.validation.models import (
     RawDataCurrentParams,
     RawDataCurrentParamsBase,
@@ -475,7 +478,7 @@ def get_osm_current_snapshot_as_file(
 
 @router.post("/snapshot/plain/")
 @version(1)
-def get_osm_current_snapshot_as_plain_geojson(
+async def get_osm_current_snapshot_as_plain_geojson(
     request: Request,
     params: RawDataCurrentParamsBase,
     user: AuthUser = Depends(get_optional_user),
@@ -484,34 +487,61 @@ def get_osm_current_snapshot_as_plain_geojson(
 
     Args:
         request (Request): _description_
-        params (RawDataCurrentParamsBase): Same as /snapshot excpet multiple output format options and configurations
+        params (RawDataCurrentParamsBase): Same as /snapshot except multiple output format options and configurations
 
     Returns:
-        Featurecollection: Geojson
+        FeatureCollection: Geojson
     """
     if user.id == 0 and params.include_user_metadata:
         raise HTTPException(
             status_code=403,
             detail=[
                 {
-                    "msg": "Insufficient Permission for extracting exports with user metadata , Please login first"
+                    "msg": "Insufficient Permission for extracting exports with user metadata, Please login first"
                 }
             ],
         )
     area_m2 = area(json.loads(params.geometry.model_dump_json()))
+
     area_km2 = area_m2 * 1e-6
-    if area_km2 > 5:
+    if int(area_km2) > 6:
         raise HTTPException(
             status_code=400,
             detail=[
                 {
-                    "msg": f"""Polygon Area {int(area_km2)} Sq.KM is higher than Threshold : 10 Sq.KM"""
+                    "msg": f"""Polygon Area {int(area_km2)} Sq.KM is higher than Threshold : 6 Sq.KM"""
                 }
             ],
         )
+
     params.output_type = "geojson"  # always geojson
-    result = RawData(params).extract_plain_geojson()
-    return result
+
+    async def generate_geojson() -> AsyncGenerator[bytes, None]:
+        # start of featurecollection
+        yield b'{"type": "FeatureCollection", "features": ['
+
+        raw_data = RawData(params)
+        extraction_query = raw_currentdata_extraction_query(params)
+
+        with raw_data.con.cursor(name="fetch_raw_quick") as cursor:
+            cursor.itersize = 500
+            cursor.execute(extraction_query)
+
+            first_feature = True
+            for row in cursor:
+                feature = orjson.loads(row[0])
+                if not first_feature:
+                    # add comma to maintain the struct
+                    yield b","
+                else:
+                    first_feature = False
+                yield orjson.dumps(feature)
+            cursor.close()
+
+        # end of featurecollect
+        yield b"]}"
+
+    return StreamingResponse(generate_geojson(), media_type="application/geo+json")
 
 
 @router.get("/countries/")
diff --git a/src/app.py b/src/app.py
index 2788efb2..a00b1e72 100644
--- a/src/app.py
+++ b/src/app.py
@@ -1913,7 +1913,7 @@ def upload_dataset(self, dump_config_to_s3=False):
                 dataset_info["hdx_upload"] = "SUCCESS"
             except Exception as ex:
                 logging.error(ex)
-                raise ex
+                # raise ex
                 dataset_info["hdx_upload"] = "FAILED"
 
             dataset_info["name"] = self.dataset["name"]