fix: csv export without tmp file

rti · rti · commit 347698fc5530 · 2025-09-04T04:32:28.000Z
diff --git a/app.py b/app.py
@@ -48,9 +48,7 @@ def read_root():
 
 
 @app.get("/csv/metrics")
-async def metric_csv(
-    background_tasks: BackgroundTasks, authorization: Optional[str] = None
-):
+async def metric_csv(authorization: Optional[str] = None):
     """Quantity CSV"""
 
     try:
@@ -63,4 +61,4 @@ async def metric_csv(
     except AssertionError:
         return PlainTextResponse("Authorization Failed", 403)
 
-    return await export_metric_csv(background_tasks)
+    return await export_metric_csv()
diff --git a/export_csv/metric.py b/export_csv/metric.py
@@ -17,7 +17,7 @@
 from model.enum import WikibaseType, WikibaseURLType
 
 
-async def export_metric_csv(background_tasks: BackgroundTasks) -> StreamingResponse:
+async def export_metric_csv() -> StreamingResponse:
     """CSV with Requested Metrics"""
 
     query = get_metrics_query()
@@ -26,7 +26,6 @@ async def export_metric_csv(background_tasks: BackgroundTasks) -> StreamingRespo
         query=query,
         export_filename="metrics",
         index_col="wikibase_id",
-        background_tasks=background_tasks,
     )
 
 
diff --git a/export_csv/util.py b/export_csv/util.py
@@ -1,8 +1,6 @@
 """Utilities"""
 
-import os
 from typing import Optional
-import uuid
 from fastapi import BackgroundTasks
 from fastapi.responses import StreamingResponse
 import pandas
@@ -25,14 +23,9 @@ async def read_sql_query(
         return df
 
 
-CHUNK_SIZE = 1024 * 1024
-EXPORT_DIRECTORY = "export/data"
-
-
 async def export_csv(
     query: Select,
     export_filename: str,
-    background_tasks: BackgroundTasks,
     index_col: Optional[str] = None,
 ):
     """Export CSV"""
@@ -41,18 +34,8 @@ async def export_csv(
     if index_col == "wikibase_id":
         assert len(set(df.index)) == len(df), "Returned Multiple Rows per Wikibase"
 
-    os.makedirs(EXPORT_DIRECTORY, exist_ok=True)
-
-    filename = f"{EXPORT_DIRECTORY}/{uuid.uuid4()}.csv"
-    df.to_csv(filename)
+    csv = df.to_csv()
     del df
 
-    def iterfile():
-        with open(filename, "rb") as f:
-            while chunk := f.read(CHUNK_SIZE):
-                yield chunk
-
-    background_tasks.add_task(os.remove, filename)
-
     headers = {"Content-Disposition": f'attachment; filename="{export_filename}.csv"'}
-    return StreamingResponse(iterfile(), headers=headers, media_type="text/csv")
+    return StreamingResponse(csv, headers=headers, media_type="text/csv")
diff --git a/tests/test_export_csv/test_export_metrics_csv.py b/tests/test_export_csv/test_export_metrics_csv.py
@@ -1,27 +1,11 @@
 """Test Export Metrics CSV"""
 
-import os
 import re
-from typing import Callable
 import pytest
-from export_csv.metric import export_metric_csv
+from fastapi.testclient import TestClient
+from app import app
 
 
-class MockBackgroundTasks:
-    """Mock BackgroundTasks"""
-
-    task_list = []
-
-    # pylint: disable-next=unused-argument
-    def add_task(self, func: Callable, *args):
-        """
-        Add Task
-
-        really add filename to list
-        """
-
-        self.task_list.append(args[0])
-
 
 EXPECTED_HEADER_ROW = ",".join(
     [
@@ -46,7 +30,7 @@ def add_task(self, func: Callable, *args):
         "bot_change_user_count",
         "software_version_observation_date",
         "software_name",
-        "version\n",
+        "version",
     ]
 )
 EXPECTED_PATTERN = re.compile(
@@ -76,7 +60,7 @@ def add_task(self, func: Callable, *args):
             # # Software
             r"(\d{4}-\d\d-\d\d \d\d:\d\d:\d\d|)",
             r"(MediaWiki|)",
-            r"(\d+\.\d+\.\d+|)\n",
+            r"(\d+\.\d+\.\d+|)",
         ]
     )
 )
@@ -94,19 +78,14 @@ def add_task(self, func: Callable, *args):
 async def test_export_metric_csv():
     """Test Export Metric CSV"""
 
-    mock_background_tasks = MockBackgroundTasks()
-
-    result = await export_metric_csv(mock_background_tasks)
-    assert len(mock_background_tasks.task_list) == 1
+    client = TestClient(app)
+    result = client.get("/csv/metrics?authorization=test-auth-token")
     assert result.status_code == 200
-    assert result.media_type == "text/csv"
-    # CANNOT FIGURE OUT HOW TO CHECK CONTENT OF RESPONSE
+    content = result.content.decode("utf-8")
 
-    with open(mock_background_tasks.task_list[0], mode="r", encoding="utf-8") as file:
-        returned_lines = file.readlines()
-        assert returned_lines[0] == EXPECTED_HEADER_ROW
-        for returned_line in returned_lines[1:]:
-            assert EXPECTED_PATTERN.match(returned_line)
+    lines = content.splitlines()
+    assert len(lines) >= 2
+    assert lines[0] == EXPECTED_HEADER_ROW
 
-    for file in mock_background_tasks.task_list:
-        os.remove(file)
+    for line in lines[1:]:
+        assert EXPECTED_PATTERN.match(line)
