diff --git a/querybook/server/lib/export/base_exporter.py b/querybook/server/lib/export/base_exporter.py
index b5b66f2c9..44b17fa33 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, TypedDict
from app.db import with_session
from env import QuerybookSettings
from lib.logger import get_logger
@@ -13,6 +13,15 @@
LOG = get_logger(__file__)
+class ExporterUrlResult(TypedDict):
+ url: str
+ message: str
+ message_type: str
+
+
+ExporterResult = Union[str, ExporterUrlResult]
+
+
class BaseExporter(metaclass=ABCMeta):
@property
@abstractmethod
@@ -59,7 +68,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 +79,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