diff --git a/querybook/server/datasources/query_execution.py b/querybook/server/datasources/query_execution.py index 8c3515cdb..929df8296 100644 --- a/querybook/server/datasources/query_execution.py +++ b/querybook/server/datasources/query_execution.py @@ -57,6 +57,9 @@ from env import QuerybookSettings from lib.notify.utils import notify_user +from pandas import DataFrame, ExcelWriter +from io import BytesIO + QUERY_RESULT_LIMIT_CONFIG = get_config_value("query_result_limit") @@ -289,11 +292,64 @@ def download_statement_execution_result(statement_execution_id): # We read the raw file and download it for the user reader.start() raw = reader.read_raw() - response = Response(raw) - response.headers["Content-Type"] = "text/csv" - response.headers[ - "Content-Disposition" - ] = f'attachment; filename="{download_file_name}"' + response = Response( + raw, + mimetype="text/csv", + headers={ + "Content-disposition": f"attachment; filename={download_file_name}" + }, + ) + return response + + +@register( + "/statement_execution//result/download_xlsx/", + methods=["GET"], + require_auth=True, + custom_response=True, +) +def download_statement_execution_result_xlsx(statement_execution_id): + with DBSession() as session: + statement_execution = logic.get_statement_execution_by_id( + statement_execution_id, session=session + ) + api_assert( + statement_execution is not None, message="Invalid statement execution" + ) + verify_query_execution_permission( + statement_execution.query_execution_id, session=session + ) + + download_file_name = f"result_{statement_execution.query_execution_id}_{statement_execution_id}.xlsx" + + reader = GenericReader(statement_execution.result_path) + reader.start() + + # Read all rows and create a DataFrame + data = reader.read_csv(number_of_lines=None) + column_names = data[0] + data_rows = data[1:] + + df = DataFrame(data_rows, columns=column_names) + + # Write the DataFrame to an Excel spreadsheet + xlsx_output = BytesIO() + with ExcelWriter(xlsx_output) as writer: + df.to_excel( + writer, + sheet_name=f"Execution {statement_execution.query_execution_id}", + index=False, + ) + data = xlsx_output.getvalue() + + # Create a Flask response with the Excel file + response = Response( + data, + mimetype="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + headers={ + "Content-disposition": f"attachment; filename={download_file_name}" + }, + ) return response diff --git a/querybook/webapp/components/StatementExecutionBar/ResultExportDropdown.tsx b/querybook/webapp/components/StatementExecutionBar/ResultExportDropdown.tsx index 619de739d..317a235fc 100644 --- a/querybook/webapp/components/StatementExecutionBar/ResultExportDropdown.tsx +++ b/querybook/webapp/components/StatementExecutionBar/ResultExportDropdown.tsx @@ -9,7 +9,10 @@ import { IStatementExecution, IStatementResult, } from 'const/queryExecution'; -import { getStatementExecutionResultDownloadUrl } from 'lib/query-execution'; +import { + getStatementExecutionResultDownloadUrl, + getStatementExecutionResultDownloadUrlXlsx, +} from 'lib/query-execution'; import { getExporterAuthentication, pollExporterTaskPromise, @@ -101,6 +104,15 @@ export const ResultExportDropdown: React.FunctionComponent = ({ } }, [statementId]); + const onDownloadXlsxClick = React.useCallback(() => { + const url = getStatementExecutionResultDownloadUrlXlsx(statementId); + if (url) { + Utils.download(url, `${statementId}.xlsx`); + } else { + toast.error('No valid url!'); + } + }, [statementId]); + const onExportTSVClick = React.useCallback(async () => { const rawResult = statementResult?.data || []; const parsedResult = tableToTSV(rawResult); @@ -179,6 +191,11 @@ export const ResultExportDropdown: React.FunctionComponent = ({ onClick: onDownloadClick, icon: 'Download', }, + { + name: 'Download Full Result (as XLSX)', + onClick: onDownloadXlsxClick, + icon: 'Sheet', + }, { name: ( diff --git a/querybook/webapp/lib/query-execution.ts b/querybook/webapp/lib/query-execution.ts index 55cbae723..bed1d1f54 100644 --- a/querybook/webapp/lib/query-execution.ts +++ b/querybook/webapp/lib/query-execution.ts @@ -1,3 +1,6 @@ export function getStatementExecutionResultDownloadUrl(id: number) { return `${location.protocol}//${location.host}/ds/statement_execution/${id}/result/download/`; } +export function getStatementExecutionResultDownloadUrlXlsx(id: number) { + return `${location.protocol}//${location.host}/ds/statement_execution/${id}/result/download_xlsx/`; +} diff --git a/querybook/webapp/ui/Icon/LucideIcons.ts b/querybook/webapp/ui/Icon/LucideIcons.ts index bac6f3be9..6bea8e0ae 100644 --- a/querybook/webapp/ui/Icon/LucideIcons.ts +++ b/querybook/webapp/ui/Icon/LucideIcons.ts @@ -93,6 +93,7 @@ import { Server, Settings, Share, + Sheet, Sidebar, Slash, Sliders, @@ -210,6 +211,7 @@ const AllLucideIcons = { Settings, Sidebar, Share, + Sheet, Slash, Sliders, Smile, diff --git a/requirements/base.txt b/requirements/base.txt index 096d94648..67a958109 100644 --- a/requirements/base.txt +++ b/requirements/base.txt @@ -46,3 +46,6 @@ numpy>=1.22.2 # not directly required, pinned by Snyk to avoid a vulnerability # Query engine - PostgreSQL psycopg2==2.9.5 + +# Excel export +openpyxl==3.1.2