Skip to content

Commit

Permalink
Dh-5004/adding the create csv endpoint without S3
Browse files Browse the repository at this point in the history
  • Loading branch information
MohammadrezaPourreza committed Nov 15, 2023
1 parent c70a26c commit ca7a777
Show file tree
Hide file tree
Showing 3 changed files with 46 additions and 0 deletions.
6 changes: 6 additions & 0 deletions dataherald/api/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,12 @@ def get_response_file(
) -> FileResponse:
pass

@abstractmethod
def create_response_csv(
self, response_id: str, background_tasks: BackgroundTasks
) -> FileResponse:
pass

@abstractmethod
def delete_golden_record(self, golden_record_id: str) -> dict:
pass
Expand Down
27 changes: 27 additions & 0 deletions dataherald/api/fastapi.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import csv
import json
import logging
import os
Expand Down Expand Up @@ -409,6 +410,32 @@ def get_response_file(
media_type="text/csv",
)

@override
def create_response_csv(
self, response_id: str, background_tasks: BackgroundTasks
) -> FileResponse:
response_repository = ResponseRepository(self.storage)
try:
response = response_repository.find_by_id(response_id)
except InvalidId as e:
raise HTTPException(status_code=400, detail=str(e)) from e

if not response:
raise HTTPException(
status_code=404, detail="Question, response, or db_connection not found"
)
csv_file_path = "query_result.csv"
with open(csv_file_path, mode="w", newline="", encoding="utf-8") as file:
writer = csv.DictWriter(file, fieldnames=response.sql_query_result.columns)
writer.writeheader()
for row in response.sql_query_result.rows:
writer.writerow(row)
response = FileResponse(
path=csv_file_path, filename="query_result.csv", media_type="text/csv"
)
background_tasks.add_task(delete_file, csv_file_path)
return response

@override
def update_response(self, response_id: str) -> Response:
response_repository = ResponseRepository(self.storage)
Expand Down
13 changes: 13 additions & 0 deletions dataherald/server/fastapi/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,13 @@ def __init__(self, settings: Settings):
tags=["Responses"],
)

self.router.add_api_route(
"/api/v1/responses/{response_id}/csv",
self.create_response_csv,
methods=["GET"],
tags=["Responses"],
)

self.router.add_api_route(
"/api/v1/responses/{response_id}",
self.update_response,
Expand Down Expand Up @@ -313,6 +320,12 @@ def get_response_file(
"""Get a response file"""
return self._api.get_response_file(response_id, background_tasks)

def create_response_csv(
self, response_id: str, background_tasks: BackgroundTasks
) -> FileResponse:
"""Get a response file"""
return self._api.create_response_csv(response_id, background_tasks)

def execute_sql_query(self, query: Query) -> tuple[str, dict]:
"""Executes a query on the given db_connection_id"""
return self._api.execute_sql_query(query)
Expand Down

0 comments on commit ca7a777

Please sign in to comment.