From 892d5c1189a8cc675267aa33517e9f080011fc4e Mon Sep 17 00:00:00 2001 From: rchandna_expedia Date: Wed, 3 Apr 2024 14:03:15 -0500 Subject: [PATCH 1/3] fix: add bom to downloaded csv --- querybook/server/datasources/query_execution.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/querybook/server/datasources/query_execution.py b/querybook/server/datasources/query_execution.py index 8c3515cdb..120752305 100644 --- a/querybook/server/datasources/query_execution.py +++ b/querybook/server/datasources/query_execution.py @@ -289,7 +289,12 @@ 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) + + # Add Byte-Order-Mark to the utf8 file to make Excel happy + raw_bytes = raw.encode() # get bytes instead of string + BOM = b"\xEF\xBB\xBF" + raw_with_bom = BOM + raw_bytes + response = Response(raw_with_bom) response.headers["Content-Type"] = "text/csv" response.headers[ "Content-Disposition" From f54ffd13418c80e419827c7a4bcb09e83321bbe2 Mon Sep 17 00:00:00 2001 From: rchandna_expedia Date: Tue, 4 Jun 2024 09:52:34 -0500 Subject: [PATCH 2/3] fix: use ExcelWriter instead of bom --- .../server/datasources/query_execution.py | 69 ++++++++++++++++--- .../ResultExportDropdown.tsx | 16 ++++- querybook/webapp/lib/query-execution.ts | 3 + querybook/webapp/ui/Icon/LucideIcons.ts | 2 + requirements/base.txt | 3 + 5 files changed, 83 insertions(+), 10 deletions(-) diff --git a/querybook/server/datasources/query_execution.py b/querybook/server/datasources/query_execution.py index 120752305..83f5331e2 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,16 +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, + mimetype="text/csv", + headers={ + "Content-disposition": f"attachment; filename={download_file_name}" + }, + ) + return response + - # Add Byte-Order-Mark to the utf8 file to make Excel happy - raw_bytes = raw.encode() # get bytes instead of string - BOM = b"\xEF\xBB\xBF" - raw_with_bom = BOM + raw_bytes - response = Response(raw_with_bom) - response.headers["Content-Type"] = "text/csv" - response.headers[ - "Content-Disposition" - ] = f'attachment; filename="{download_file_name}"' +@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..00e1419f3 100644 --- a/querybook/webapp/components/StatementExecutionBar/ResultExportDropdown.tsx +++ b/querybook/webapp/components/StatementExecutionBar/ResultExportDropdown.tsx @@ -9,7 +9,7 @@ import { IStatementExecution, IStatementResult, } from 'const/queryExecution'; -import { getStatementExecutionResultDownloadUrl } from 'lib/query-execution'; +import { getStatementExecutionResultDownloadUrl, getStatementExecutionResultDownloadUrlXlsx } from 'lib/query-execution'; import { getExporterAuthentication, pollExporterTaskPromise, @@ -101,6 +101,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 +188,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..567e108d0 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 \ No newline at end of file From 577ad7a3ecfaced0aae2b7d45603387c2c373c82 Mon Sep 17 00:00:00 2001 From: rchandna_expedia Date: Tue, 4 Jun 2024 11:33:19 -0500 Subject: [PATCH 3/3] fix: lint --- querybook/server/datasources/query_execution.py | 2 +- .../StatementExecutionBar/ResultExportDropdown.tsx | 9 ++++++--- requirements/base.txt | 2 +- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/querybook/server/datasources/query_execution.py b/querybook/server/datasources/query_execution.py index 83f5331e2..929df8296 100644 --- a/querybook/server/datasources/query_execution.py +++ b/querybook/server/datasources/query_execution.py @@ -324,7 +324,7 @@ def download_statement_execution_result_xlsx(statement_execution_id): 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] diff --git a/querybook/webapp/components/StatementExecutionBar/ResultExportDropdown.tsx b/querybook/webapp/components/StatementExecutionBar/ResultExportDropdown.tsx index 00e1419f3..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, getStatementExecutionResultDownloadUrlXlsx } from 'lib/query-execution'; +import { + getStatementExecutionResultDownloadUrl, + getStatementExecutionResultDownloadUrlXlsx, +} from 'lib/query-execution'; import { getExporterAuthentication, pollExporterTaskPromise, @@ -108,7 +111,7 @@ export const ResultExportDropdown: React.FunctionComponent = ({ } else { toast.error('No valid url!'); } - }, [statementId]) + }, [statementId]); const onExportTSVClick = React.useCallback(async () => { const rawResult = statementResult?.data || []; @@ -191,7 +194,7 @@ export const ResultExportDropdown: React.FunctionComponent = ({ { name: 'Download Full Result (as XLSX)', onClick: onDownloadXlsxClick, - icon: 'Sheet' + icon: 'Sheet', }, { name: ( diff --git a/requirements/base.txt b/requirements/base.txt index 567e108d0..67a958109 100644 --- a/requirements/base.txt +++ b/requirements/base.txt @@ -48,4 +48,4 @@ numpy>=1.22.2 # not directly required, pinned by Snyk to avoid a vulnerability psycopg2==2.9.5 # Excel export -openpyxl==3.1.2 \ No newline at end of file +openpyxl==3.1.2