fix: add bom to downloaded csv for excel

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/<int:statement_execution_id>/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
 
 

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<IProps> = ({
         }
     }, [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<IProps> = ({
                 onClick: onDownloadClick,
                 icon: 'Download',
             },
+            {
+                name: 'Download Full Result (as XLSX)',
+                onClick: onDownloadXlsxClick,
+                icon: 'Sheet',
+            },
             {
                 name: (
                     <span>

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/`;
+}

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,

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