feat: show messages for URL result exporters

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()
 

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

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})

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 }) => (
-    <Modal onHide={onHide} title={title}>
-        <div className="flex-center mv24">
-            <Button title="View Export" onClick={() => window.open(url)} />
-        </div>
-    </Modal>
-);
+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 (
+        <Modal onHide={onHide} title="Export Complete">
+            {message ? <Message type={messageType}>{message}</Message> : null}
+            <div className="flex-center mv24">
+                <Button title="View Export" onClick={() => open(url)} />
+            </div>
+        </Modal>
+    );
+};
 
 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' ? (
-        <UrlModal
-            url={exportedInfo.info}
-            title={title}
+        <UrlExportRedirect
+            exportResult={
+                typeof exportedInfo.info === 'string'
+                    ? { url: exportedInfo.info }
+                    : exportedInfo.info
+            }
             onHide={() => toast.dismiss(taskId)}
         />
     ) : (
         <CopyPasteModal
-            text={exportedInfo.info}
+            text={String(exportedInfo.info)}
             title={title}
             onHide={() => toast.dismiss(taskId)}
         />

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