From 5b543e10b3a142564c2f8772d6d1d73cdd296b99 Mon Sep 17 00:00:00 2001
From: Biel Stela <biel.stela@gmail.com>
Date: Thu, 12 Sep 2024 12:20:08 +0200
Subject: [PATCH] Tidy up the table response model

---
 api/app/models/grid.py  |  9 +++++++++
 api/app/routers/grid.py |  7 +++----
 api/tests/test_grid.py  | 25 +++++++++++--------------
 3 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/api/app/models/grid.py b/api/app/models/grid.py
index f1ea4b3f..17ce102d 100644
--- a/api/app/models/grid.py
+++ b/api/app/models/grid.py
@@ -136,3 +136,12 @@ def to_sql_query(self, table_name: str) -> str:
             )
         )
         return str(query.compile(compile_kwargs={"literal_binds": True}))
+
+
+class TableResultColumn(BaseModel):
+    column: Annotated[str, Field(title="column", description="Column name")]
+    values: Annotated[list, Field(description="Check dataset metadata for type info")]
+
+
+class TableResults(BaseModel):
+    table: list[TableResultColumn]
diff --git a/api/app/routers/grid.py b/api/app/routers/grid.py
index 3012283e..c4e84e40 100644
--- a/api/app/routers/grid.py
+++ b/api/app/routers/grid.py
@@ -10,7 +10,6 @@
 import shapely
 from fastapi import APIRouter, Depends, HTTPException, Path, Query
 from fastapi.params import Body
-from fastapi.responses import ORJSONResponse
 from geojson_pydantic import Feature
 from h3 import H3CellError
 from h3ronpy.polars import cells_to_string
@@ -19,7 +18,7 @@
 from starlette.responses import Response
 
 from app.config.config import get_settings
-from app.models.grid import MultiDatasetMeta, TableFilters
+from app.models.grid import MultiDatasetMeta, TableFilters, TableResults
 
 log = logging.getLogger("uvicorn.error")
 
@@ -118,7 +117,7 @@ def read_table(
     level: Annotated[int, Query(..., description="Tile level at which the query will be computed")],
     filters: TableFilters = Depends(),
     geojson: Annotated[Feature | None, Body(description="GeoJSON feature used to filter the cells.")] = None,
-) -> ORJSONResponse:
+) -> TableResults:
     """Query tile dataset and return table data"""
     files_path = pathlib.Path(get_settings().grid_tiles_path) / str(level)
     if not files_path.exists():
@@ -144,4 +143,4 @@ def read_table(
         log.exception(e)
         raise HTTPException(status_code=422, detail=str(e)) from None
 
-    return ORJSONResponse(res.to_dict(as_series=False))
+    return TableResults(table=[{"column": k, "values": v} for k, v in res.to_dict(as_series=False).items()])
diff --git a/api/tests/test_grid.py b/api/tests/test_grid.py
index fb7726d4..b75c254c 100644
--- a/api/tests/test_grid.py
+++ b/api/tests/test_grid.py
@@ -221,12 +221,11 @@ def test_grid_table(grid_dataset):
     response = test_client.post("/grid/table?level=4&order_by=-population", headers=HEADERS, content=json.dumps(body))
     assert response.status_code == 200
     assert json.loads(response.read()) == {
-        "cell": [
-            "865f00007ffffff",
-            "895f4261e03ffff",
-        ],
-        "landcover": [4, 1],
-        "population": [200, 100],
+        "table": [
+            {"column": "cell", "values": ["865f00007ffffff", "895f4261e03ffff"]},
+            {"column": "landcover", "values": [4, 1]},
+            {"column": "population", "values": [200, 100]},
+        ]
     }
 
 
@@ -241,11 +240,11 @@ def test_grid_table_geojson(grid_dataset, geojson):
     response = test_client.post("/grid/table?level=4&order_by=-population", headers=HEADERS, content=json.dumps(body))
     assert response.status_code == 200
     assert json.loads(response.read()) == {
-        "cell": [
-            "895f4261e03ffff",
-        ],
-        "landcover": [1],
-        "population": [100],
+        "table": [
+            {"column": "cell", "values": ["895f4261e03ffff"]},
+            {"column": "landcover", "values": [1]},
+            {"column": "population", "values": [100]},
+        ]
     }
 
 
@@ -258,9 +257,7 @@ def test_grid_tile_post_geojson(grid_dataset, geojson):
     )
     assert response.status_code == 200
     assert pl.read_ipc(response.read()).to_dict(as_series=False) == {
-        "cell": [
-            "895f4261e03ffff",
-        ],
+        "cell": ["895f4261e03ffff"],
         "landcover": [1],
         "population": [100],
     }