From eedeaf2842923ae4ac96fe67cafd97788f910caa Mon Sep 17 00:00:00 2001 From: CenturySturgeon Date: Fri, 31 Jul 2026 19:58:38 -0600 Subject: [PATCH] feat: show messages for URL result exporters --- querybook/server/lib/export/base_exporter.py | 13 +++-- .../server/tasks/export_query_execution.py | 5 +- .../test_export_query_execution.py | 37 +++++++++++++ .../ResultExportSuccessToast.tsx | 54 +++++++++++++------ querybook/webapp/const/queryExecution.ts | 10 +++- 5 files changed, 94 insertions(+), 25 deletions(-) create mode 100644 querybook/tests/test_lib/test_export/test_export_query_execution.py diff --git a/querybook/server/lib/export/base_exporter.py b/querybook/server/lib/export/base_exporter.py index b5b66f2c9..8e27d5468 100644 --- a/querybook/server/lib/export/base_exporter.py +++ b/querybook/server/lib/export/base_exporter.py @@ -1,5 +1,5 @@ from abc import ABCMeta, abstractmethod -from typing import Generator, List +from typing import Generator, List, Union from app.db import with_session from env import QuerybookSettings from lib.logger import get_logger @@ -12,6 +12,8 @@ LOG = get_logger(__file__) +ExporterResult = Union[str, dict[str, str]] + class BaseExporter(metaclass=ABCMeta): @property @@ -59,7 +61,9 @@ def export_form(self): return None @abstractmethod - def export(self, statement_execution_id: int, uid: int, **options) -> str: + def export( + self, statement_execution_id: int, uid: int, **options + ) -> ExporterResult: """This function exports the query results of statement_execution_id to given output Arguments: @@ -68,8 +72,9 @@ def export(self, statement_execution_id: int, uid: int, **options) -> str: options {[Dict]} -- optional additional options, note they must be optional since Returns: - str -- String for frontend to display - Behavior noted by EXPORTER_TYPE + str -- String for frontend to display. URL exporters may instead + return {"url": str, "message": str, "message_type": "warning"} + to show a message before redirecting. """ raise NotImplementedError() diff --git a/querybook/server/tasks/export_query_execution.py b/querybook/server/tasks/export_query_execution.py index 92a6bb0eb..77f44ff2d 100644 --- a/querybook/server/tasks/export_query_execution.py +++ b/querybook/server/tasks/export_query_execution.py @@ -1,5 +1,6 @@ from typing import Dict from app.flask_app import celery +from lib.export.base_exporter import ExporterResult from lib.export.all_exporters import get_exporter from lib.logger import get_logger @@ -17,7 +18,7 @@ def export_query_execution_task( ): exporter = get_exporter(exporter_name) try: - info = exporter.export( + info: ExporterResult = exporter.export( statement_execution_id, current_user_id, **(exporter_params or {}) ) return { @@ -26,4 +27,4 @@ def export_query_execution_task( } except Exception as e: LOG.error(e) - raise e + raise e \ No newline at end of file diff --git a/querybook/tests/test_lib/test_export/test_export_query_execution.py b/querybook/tests/test_lib/test_export/test_export_query_execution.py new file mode 100644 index 000000000..56db697ff --- /dev/null +++ b/querybook/tests/test_lib/test_export/test_export_query_execution.py @@ -0,0 +1,37 @@ +from unittest import TestCase, mock + +from tasks.export_query_execution import export_query_execution_task + + +class ExportQueryExecutionTaskTestCase(TestCase): + def setUp(self): + self.exporter = mock.Mock(exporter_type="url") + self.get_exporter_patch = mock.patch( + "tasks.export_query_execution.get_exporter", return_value=self.exporter + ) + self.get_exporter_patch.start() + self.addCleanup(self.get_exporter_patch.stop) + + def test_preserves_legacy_url_string_result(self): + self.exporter.export.return_value = "https://example.com/export" + + result = export_query_execution_task.run( + "Example exporter", 1, 2, {"format": "csv"} + ) + + self.assertEqual( + result, + {"type": "url", "info": "https://example.com/export"}, + ) + + def test_preserves_structured_url_result(self): + export_result = { + "url": "https://example.com/export", + "message": "Your export contains filtered rows.", + "message_type": "warning", + } + self.exporter.export.return_value = export_result + + result = export_query_execution_task.run("Example exporter", 1, 2, {}) + + self.assertEqual(result, {"type": "url", "info": export_result}) \ No newline at end of file diff --git a/querybook/webapp/components/StatementExecutionBar/ResultExportSuccessToast.tsx b/querybook/webapp/components/StatementExecutionBar/ResultExportSuccessToast.tsx index 5420f0d39..79bdbe6f8 100644 --- a/querybook/webapp/components/StatementExecutionBar/ResultExportSuccessToast.tsx +++ b/querybook/webapp/components/StatementExecutionBar/ResultExportSuccessToast.tsx @@ -1,38 +1,58 @@ import React from 'react'; import toast from 'react-hot-toast'; -import { IQueryExecutionExportStatusInfo } from 'const/queryExecution'; +import { + IQueryExecutionExportStatusInfo, + IUrlExportResult, +} from 'const/queryExecution'; import { Button } from 'ui/Button/Button'; import { CopyPasteModal } from 'ui/CopyPasteModal/CopyPasteModal'; +import { Message } from 'ui/Message/Message'; import { Modal } from 'ui/Modal/Modal'; -const UrlModal: React.FunctionComponent<{ - url: string; - title: string; - onHide: () => any; -}> = ({ url, title, onHide }) => ( - -
-
- -); +const openExport = (url: string) => window.open(url); + +export const UrlExportRedirect: React.FunctionComponent<{ + exportResult: IUrlExportResult; + onHide: () => void; + open?: (url: string) => void; +}> = ({ + exportResult: { url, message, message_type: messageType = 'warning' }, + onHide, + open = openExport, +}) => { + return ( + + {message ? {message} : null} +
+
+ + ); +}; export const ResultExportSuccessToast = ( - data: IQueryExecutionExportStatusInfo + data: IQueryExecutionExportStatusInfo, ) => { const { result: exportedInfo, task_id: taskId } = data; const title = 'Export Complete'; + if (!exportedInfo) { + return null; + } + return exportedInfo.type === 'url' ? ( - toast.dismiss(taskId)} /> ) : ( toast.dismiss(taskId)} /> diff --git a/querybook/webapp/const/queryExecution.ts b/querybook/webapp/const/queryExecution.ts index 791641e2d..5dd4cd08d 100644 --- a/querybook/webapp/const/queryExecution.ts +++ b/querybook/webapp/const/queryExecution.ts @@ -73,9 +73,15 @@ export interface IQueryExecution { total?: number; } +export interface IUrlExportResult { + url: string; + message?: string; + message_type?: 'warning'; +} + export interface IQueryExecutionExportResult { type: 'url' | 'text'; - info: string; + info: string | IUrlExportResult; } export interface IQueryExecutionExportStatusInfo { @@ -177,4 +183,4 @@ export interface IQueryValidationResult { severity: string; type: string; suggestion: string | null; -} +} \ No newline at end of file