Remove temp file when is downloaded

Dataherald · Nov 7, 2023 · 29112a3 · 29112a3
1 parent 4767606
commit 29112a3
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 5 deletions.
diff --git a/dataherald/api/__init__.py b/dataherald/api/__init__.py
@@ -131,7 +131,9 @@ def get_response(self, response_id: str) -> Response:
         pass
 
     @abstractmethod
-    def get_response_file(self, response_id: str) -> FileResponse:
+    def get_response_file(
+        self, response_id: str, background_tasks: BackgroundTasks
+    ) -> FileResponse:
         pass
 
     @abstractmethod

diff --git a/dataherald/api/fastapi.py b/dataherald/api/fastapi.py
@@ -65,6 +65,10 @@ def async_scanning(scanner, database, scanner_request, storage):
     )
 
 
+def delete_file(file_location: str):
+    os.remove(file_location)
+
+
 class FastAPI(API):
     def __init__(self, system: System):
         super().__init__(system)
@@ -369,7 +373,9 @@ def get_response(self, response_id: str) -> Response:
         return result
 
     @override
-    def get_response_file(self, response_id: str) -> FileResponse:
+    def get_response_file(
+        self, response_id: str, background_tasks: BackgroundTasks
+    ) -> FileResponse:
         response_repository = ResponseRepository(self.storage)
         question_repository = QuestionRepository(self.storage)
         db_connection_repository = DatabaseConnectionRepository(self.storage)
@@ -389,8 +395,11 @@ def get_response_file(self, response_id: str) -> FileResponse:
 
         s3 = S3()
 
+        file_location = s3.download(result.csv_file_path, db_connection.file_storage)
+        background_tasks.add_task(delete_file, file_location)
+
         return FileResponse(
-            s3.download(result.csv_file_path, db_connection.file_storage),
+            file_location,
             media_type="text/csv",
         )
 

diff --git a/dataherald/server/fastapi/__init__.py b/dataherald/server/fastapi/__init__.py
@@ -307,9 +307,11 @@ def update_response(self, response_id: str) -> Response:
         """Update a response"""
         return self._api.update_response(response_id)
 
-    def get_response_file(self, response_id: str) -> FileResponse:
+    def get_response_file(
+        self, response_id: str, background_tasks: BackgroundTasks
+    ) -> FileResponse:
         """Get a response file"""
-        return self._api.get_response_file(response_id)
+        return self._api.get_response_file(response_id, background_tasks)
 
     def execute_sql_query(self, query: Query) -> tuple[str, dict]:
         """Executes a query on the given db_connection_id"""