DH-4905 bug fix on csv file generation endpoint (#236)

Dataherald · Nov 6, 2023 · 682a402 · 682a402
1 parent 0d1daee
commit 682a402
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 7 deletions.
diff --git a/dataherald/api/__init__.py b/dataherald/api/__init__.py
@@ -54,6 +54,10 @@ def answer_question_with_timeout(
     ) -> Response:
         pass
 
+    @abstractmethod
+    def update_response(self, response_id: str) -> Response:
+        pass
+
     @abstractmethod
     def get_questions(self, db_connection_id: str | None = None) -> list[Question]:
         pass

diff --git a/dataherald/api/fastapi.py b/dataherald/api/fastapi.py
@@ -372,24 +372,32 @@ def get_response(self, response_id: str) -> Response:
     @override
     def get_response_file(self, response_id: str) -> JSONResponse:
         response_repository = ResponseRepository(self.storage)
-
+        question_repository = QuestionRepository(self.storage)
+        db_connection_repository = DatabaseConnectionRepository(self.storage)
         try:
             result = response_repository.find_by_id(response_id)
+            question = question_repository.find_by_id(result.question_id)
+            db_connection = db_connection_repository.find_by_id(
+                question.db_connection_id
+            )
         except InvalidId as e:
             raise HTTPException(status_code=400, detail=str(e)) from e
 
         if not result:
-            raise HTTPException(status_code=404, detail="Question not found")
+            raise HTTPException(
+                status_code=404, detail="Question, response, or db_connection not found"
+            )
 
         s3 = S3()
         return JSONResponse(
             status_code=201,
             content={
-                "csv_download_url": s3.download_url(result.csv_file_path),
+                "csv_download_url": s3.download_url(
+                    result.csv_file_path, db_connection.file_storage
+                ),
             },
         )
 
-
     @override
     def update_response(self, response_id: str) -> Response:
         response_repository = ResponseRepository(self.storage)
@@ -415,7 +423,6 @@ def update_response(self, response_id: str) -> Response:
             raise HTTPException(status_code=404, detail=str(e)) from e
         return response
 
-
     @override
     def get_questions(self, db_connection_id: str | None = None) -> list[Question]:
         question_repository = QuestionRepository(self.storage)
@@ -488,7 +495,10 @@ def create_response(
                 context = context_store.retrieve_context_for_question(user_question)
                 start_generated_answer = time.time()
                 response = sql_generation.generate_response(
-                    user_question, database_connection, context[0]
+                    user_question,
+                    database_connection,
+                    context[0],
+                    large_query_result_in_csv,
                 )
             else:
                 response = Response(

diff --git a/dataherald/server/fastapi/__init__.py b/dataherald/server/fastapi/__init__.py
@@ -240,7 +240,9 @@ def answer_question(
             return self._api.answer_question_with_timeout(
                 run_evaluator, large_query_result_in_csv, question_request
             )
-        return self._api.answer_question(run_evaluator, large_query_result_in_csv, question_request)
+        return self._api.answer_question(
+            run_evaluator, large_query_result_in_csv, question_request
+        )
 
     def get_questions(self, db_connection_id: str | None = None) -> list[Question]:
         return self._api.get_questions(db_connection_id)