feat: Added excel export

plugins/exporter_plugin/__init__.py

@@ -1,3 +1,7 @@
-ALL_PLUGIN_EXPORTERS = [
+from lib.export.exporters.python_exporter import PythonExporter
+from lib.export.exporters.r_exporter import RExporter
 
+ALL_PLUGIN_EXPORTERS = [
+    PythonExporter(),
+    RExporter()
 ]

querybook/server/datasources/query_execution.py

@@ -1,6 +1,7 @@
 from datetime import datetime
 from typing import Dict, Optional
 
+import pandas as pd
 from flask import abort, Response, redirect, request
 from flask_login import current_user
 
@@ -356,6 +357,10 @@ def get_query_execution_metadata(query_execution_id):
     custom_response=True,
 )
 def download_statement_execution_result(statement_execution_id):
+    # Get the type parameter from the request, default to csv
+    type_param = request.args.get("type", "csv")
+    LOG.info(f"Received download request for: {statement_execution_id} of type {type_param}")
+
     with DBSession() as session:
         statement_execution = logic.get_statement_execution_by_id(
             statement_execution_id, session=session
@@ -367,24 +372,62 @@ def download_statement_execution_result(statement_execution_id):
             statement_execution.query_execution_id, session=session
         )
 
-        download_file_name = f"result_{statement_execution.query_execution_id}_{statement_execution_id}.csv"
+        file_extension = "csv" if type_param == "csv" else "xlsx"
+        download_file_name = f"result_{statement_execution.query_execution_id}_{statement_execution_id}.{file_extension}"
 
         reader = GenericReader(statement_execution.result_path)
         response = None
-        if reader.has_download_url:
-            # If the Reader can generate a download,
-            # we let user download the file by redirection
-            download_url = reader.get_download_url(custom_name=download_file_name)
-            response = redirect(download_url)
-        else:
-            # 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}"'
+
+        try:
+            reader.get_download_url(custom_name=download_file_name)
+            LOG.info(f"Download URL: {reader.get_download_url(custom_name=download_file_name)}")
+        except NotImplementedError:
+            LOG.info("Download URL not implemented")
+
+        if type_param == "csv":
+            if reader.has_download_url:
+                # If the Reader can generate a download,
+                # we let user download the file by redirection
+                download_url = reader.get_download_url(custom_name=download_file_name)
+                response = redirect(download_url)
+            else:
+                # 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}"'
             )
+        else:  # Excel format
+            # Read the CSV data
+            reader.start()
+            csv_data = reader.read_csv(None)
+
+            if csv_data and len(csv_data) > 0:
+                # Convert to pandas DataFrame
+                # First row is headers, rest is data
+                headers = csv_data[0]
+                data = csv_data[1:] if len(csv_data) > 1 else []
+
+                # Create DataFrame
+                df = pd.DataFrame(data, columns=headers)
+
+                # Create Excel file in memory
+                import io
+                output = io.BytesIO()
+                with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
+                    df.to_excel(writer, sheet_name='Results', index=False)
+
+                output.seek(0)
+
+                # Create response
+                response = Response(output.getvalue())
+                response.headers["Content-Type"] = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
+                response.headers[
+                    "Content-Disposition"
+                ] = f'attachment; filename="{download_file_name}"'
         return response
 
 

querybook/webapp/components/StatementExecutionBar/ResultExportDropdown.tsx

@@ -118,6 +118,17 @@ export const ResultExportDropdown: React.FunctionComponent<IProps> = ({
         }
     }, [statementId]);
 
+    const onDownloadXlsxClick = React.useCallback(() => {
+        trackExportButtonClick('Download XLSX');
+
+        const url = getStatementExecutionResultDownloadUrl(statementId, "xlsx");
+        if (url) {
+            Utils.download(url, `${statementId}.xlsx`);
+        } else {
+            toast.error('No valid url!');
+        }
+    }, [statementId]);
+
     const onExportTSVClick = React.useCallback(async () => {
         trackExportButtonClick('Copy to Clipboard');
 
@@ -200,6 +211,11 @@ export const ResultExportDropdown: React.FunctionComponent<IProps> = ({
                 onClick: onDownloadClick,
                 icon: 'Download',
             },
+            {
+                name: 'Download Full Result (as xlsx)',
+                onClick: onDownloadXlsxClick,
+                icon: 'Download',
+            },
             {
                 name: (
                     <span>

querybook/webapp/lib/query-execution.ts

@@ -1,3 +1,3 @@
-export function getStatementExecutionResultDownloadUrl(id: number) {
-    return `${location.protocol}//${location.host}/ds/statement_execution/${id}/result/download/`;
+export function getStatementExecutionResultDownloadUrl(id: number, type: "csv" | "xlsx" = "csv"): string {
+    return `${location.protocol}//${location.host}/ds/statement_execution/${id}/result/download/?type=${type}`;
 }

requirements/base.txt

@@ -43,6 +43,7 @@ pandas==1.3.5
 typing-extensions==4.9.0
 setuptools>=65.5.1 # not directly required, pinned by Snyk to avoid a vulnerability
 numpy>=1.22.2,<2.0.0 # not directly required, pinned by Snyk to avoid a vulnerability
+xlsxwriter==3.2.2
 
 # Query engine - PostgreSQL
 psycopg2==2.9.5