diff --git a/plugins/exporter_plugin/__init__.py b/plugins/exporter_plugin/__init__.py index 21350fa30..4ac8e87b9 100644 --- a/plugins/exporter_plugin/__init__.py +++ b/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() ] diff --git a/querybook/server/datasources/query_execution.py b/querybook/server/datasources/query_execution.py index e70122b30..ce27554b6 100644 --- a/querybook/server/datasources/query_execution.py +++ b/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 diff --git a/querybook/webapp/components/StatementExecutionBar/ResultExportDropdown.tsx b/querybook/webapp/components/StatementExecutionBar/ResultExportDropdown.tsx index a59fd70ba..6da8de742 100644 --- a/querybook/webapp/components/StatementExecutionBar/ResultExportDropdown.tsx +++ b/querybook/webapp/components/StatementExecutionBar/ResultExportDropdown.tsx @@ -118,6 +118,17 @@ export const ResultExportDropdown: React.FunctionComponent = ({ } }, [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 = ({ onClick: onDownloadClick, icon: 'Download', }, + { + name: 'Download Full Result (as xlsx)', + onClick: onDownloadXlsxClick, + icon: 'Download', + }, { name: ( diff --git a/querybook/webapp/lib/query-execution.ts b/querybook/webapp/lib/query-execution.ts index 55cbae723..522c6dc41 100644 --- a/querybook/webapp/lib/query-execution.ts +++ b/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}`; } diff --git a/requirements/base.txt b/requirements/base.txt index 9c845b6f9..4a4e2137a 100644 --- a/requirements/base.txt +++ b/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