diff --git a/superset-frontend/src/explore/components/ExploreChartPanel.jsx b/superset-frontend/src/explore/components/ExploreChartPanel.jsx index 3bbd1512a7860..2de2d297541d2 100644 --- a/superset-frontend/src/explore/components/ExploreChartPanel.jsx +++ b/superset-frontend/src/explore/components/ExploreChartPanel.jsx @@ -19,7 +19,7 @@ import React, { useState, useEffect, useCallback, useMemo } from 'react'; import PropTypes from 'prop-types'; import Split from 'react-split'; -import { styled, useTheme } from '@superset-ui/core'; +import { styled, SupersetClient, useTheme } from '@superset-ui/core'; import { useResizeDetector } from 'react-resize-detector'; import { chartPropShape } from 'src/dashboard/util/propShapes'; import ChartContainer from 'src/chart/ChartContainer'; @@ -29,6 +29,7 @@ import { } from 'src/utils/localStorageHelpers'; import ConnectedExploreChartHeader from './ExploreChartHeader'; import { DataTablesPane } from './DataTablesPane'; +import { buildV1ChartDataPayload } from '../exploreUtils'; const propTypes = { actions: PropTypes.object.isRequired, @@ -128,6 +129,34 @@ const ExploreChartPanel = props => { getFromLocalStorage(STORAGE_KEYS.sizes, INITIAL_SIZES), ); + const { slice } = props; + const updateQueryContext = useCallback( + async function fetchChartData() { + if (slice && slice.query_context === null) { + const queryContext = buildV1ChartDataPayload({ + formData: slice.form_data, + force: false, + resultFormat: 'json', + resultType: 'full', + setDataMask: null, + ownState: null, + }); + + await SupersetClient.put({ + endpoint: `/api/v1/chart/${slice.slice_id}`, + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + query_context: JSON.stringify(queryContext), + }), + }); + } + }, + [slice], + ); + useEffect(() => { + updateQueryContext(); + }, [updateQueryContext]); + const calcSectionHeight = useCallback( percent => { let headerHeight; diff --git a/superset-frontend/src/explore/components/PropertiesModal/index.tsx b/superset-frontend/src/explore/components/PropertiesModal/index.tsx index b1201e7b34e9d..0c46f6ba792f1 100644 --- a/superset-frontend/src/explore/components/PropertiesModal/index.tsx +++ b/superset-frontend/src/explore/components/PropertiesModal/index.tsx @@ -23,11 +23,10 @@ import Button from 'src/components/Button'; import { OptionsType } from 'react-select/src/types'; import { AsyncSelect } from 'src/components/Select'; import rison from 'rison'; -import { t, SupersetClient, QueryFormData } from '@superset-ui/core'; +import { t, SupersetClient } from '@superset-ui/core'; import Chart, { Slice } from 'src/types/Chart'; import { Form, FormItem } from 'src/components/Form'; import { getClientErrorObject } from 'src/utils/getClientErrorObject'; -import { buildV1ChartDataPayload } from '../../exploreUtils'; type PropertiesModalProps = { slice: Slice; @@ -82,26 +81,6 @@ export default function PropertiesModal({ label: `${owner.first_name} ${owner.last_name}`, })), ); - - if (chart.query_context === null) { - // set query_context if null - const queryContext = buildV1ChartDataPayload({ - formData: slice.form_data as QueryFormData, - force: false, - resultFormat: 'json', - resultType: 'full', - setDataMask: null, - ownState: null, - }); - - await SupersetClient.put({ - endpoint: `/api/v1/chart/${slice.slice_id}`, - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - query_context: JSON.stringify(queryContext), - }), - }); - } } catch (response) { const clientError = await getClientErrorObject(response); showError(clientError); diff --git a/superset-frontend/src/types/Chart.ts b/superset-frontend/src/types/Chart.ts index 02b9025c83874..4c670482b43b2 100644 --- a/superset-frontend/src/types/Chart.ts +++ b/superset-frontend/src/types/Chart.ts @@ -47,6 +47,7 @@ export type Slice = { description: string | null; cache_timeout: number | null; form_data?: QueryFormData; + query_context?: object; }; export default Chart; diff --git a/superset/models/slice.py b/superset/models/slice.py index cf6addbb829f2..bc37cc39859e9 100644 --- a/superset/models/slice.py +++ b/superset/models/slice.py @@ -193,6 +193,7 @@ def data(self) -> Dict[str, Any]: "description_markeddown": self.description_markeddown, "edit_url": self.edit_url, "form_data": self.form_data, + "query_context": self.query_context, "modified": self.modified(), "owners": [ f"{owner.first_name} {owner.last_name}" for owner in self.owners diff --git a/superset/reports/commands/execute.py b/superset/reports/commands/execute.py index 2ce1346ed66bc..070223e0753df 100644 --- a/superset/reports/commands/execute.py +++ b/superset/reports/commands/execute.py @@ -207,11 +207,29 @@ def _get_screenshot(self) -> bytes: return image_data def _get_csv_data(self) -> bytes: - if self._report_schedule.chart: - url = self._get_url(csv=True) - auth_cookies = machine_auth_provider_factory.instance.get_auth_cookies( - self._get_user() - ) + url = self._get_url(csv=True) + auth_cookies = machine_auth_provider_factory.instance.get_auth_cookies( + self._get_user() + ) + + # To load CSV data from the endpoint the chart must have been saved + # with its query context. For charts without saved query context we + # get a screenshot to force the chart to produce and save the query + # context. + if self._report_schedule.chart.query_context is None: + logger.warning("No query context found, taking a screenshot to generate it") + try: + self._get_screenshot() + except ( + ReportScheduleScreenshotFailedError, + ReportScheduleScreenshotTimeout, + ) as ex: + raise ReportScheduleCsvFailedError( + "Unable to fetch CSV data because the chart has no query context " + "saved, and an error occurred when fetching it via a screenshot. " + "Please try loading the chart and saving it again." + ) from ex + try: csv_data = get_chart_csv_data(url, auth_cookies) except SoftTimeLimitExceeded: diff --git a/tests/integration_tests/reports/commands_tests.py b/tests/integration_tests/reports/commands_tests.py index 926d8c513e8f0..b0b4e622c6dbd 100644 --- a/tests/integration_tests/reports/commands_tests.py +++ b/tests/integration_tests/reports/commands_tests.py @@ -50,7 +50,6 @@ ReportScheduleNotFoundError, ReportScheduleNotificationError, ReportSchedulePreviousWorkingError, - ReportSchedulePruneLogError, ReportScheduleScreenshotFailedError, ReportScheduleScreenshotTimeout, ReportScheduleWorkingTimeoutError, @@ -133,6 +132,7 @@ def create_report_notification( validator_config_json: Optional[str] = None, grace_period: Optional[int] = None, report_format: Optional[ReportDataFormat] = None, + name: Optional[str] = None, ) -> ReportSchedule: report_type = report_type or ReportScheduleType.REPORT target = email_target or slack_channel @@ -154,11 +154,14 @@ def create_report_notification( recipient_config_json=json.dumps(config_json), ) + if name is None: + name = "report_with_csv" if report_format else "report" + report_schedule = insert_report_schedule( type=report_type, - name=f"report_with_csv" if report_format else f"report", - crontab=f"0 9 * * *", - description=f"Daily report", + name=name, + crontab="0 9 * * *", + description="Daily report", sql=sql, chart=chart, dashboard=dashboard, @@ -217,6 +220,7 @@ def create_report_email_chart(): def create_report_email_chart_with_csv(): with app.app_context(): chart = db.session.query(Slice).first() + chart.query_context = '{"mock": "query_context"}' report_schedule = create_report_notification( email_target="target@email.com", chart=chart, @@ -226,6 +230,21 @@ def create_report_email_chart_with_csv(): cleanup_report_schedule(report_schedule) +@pytest.fixture() +def create_report_email_chart_with_csv_no_query_context(): + with app.app_context(): + chart = db.session.query(Slice).first() + chart.query_context = None + report_schedule = create_report_notification( + email_target="target@email.com", + chart=chart, + report_format=ReportDataFormat.DATA, + name="report_csv_no_query_context", + ) + yield report_schedule + cleanup_report_schedule(report_schedule) + + @pytest.fixture() def create_report_email_dashboard(): with app.app_context(): @@ -254,6 +273,7 @@ def create_report_slack_chart(): def create_report_slack_chart_with_csv(): with app.app_context(): chart = db.session.query(Slice).first() + chart.query_context = '{"mock": "query_context"}' report_schedule = create_report_notification( slack_channel="slack_channel", chart=chart, @@ -660,6 +680,47 @@ def test_email_chart_report_schedule_with_csv( assert_log(ReportState.SUCCESS) +@pytest.mark.usefixtures( + "load_birth_names_dashboard_with_slices", + "create_report_email_chart_with_csv_no_query_context", +) +@patch("superset.utils.csv.urllib.request.urlopen") +@patch("superset.utils.csv.urllib.request.OpenerDirector.open") +@patch("superset.reports.notifications.email.send_email_smtp") +@patch("superset.utils.csv.get_chart_csv_data") +@patch("superset.utils.screenshots.ChartScreenshot.get_screenshot") +def test_email_chart_report_schedule_with_csv_no_query_context( + screenshot_mock, + csv_mock, + email_mock, + mock_open, + mock_urlopen, + create_report_email_chart_with_csv_no_query_context, +): + """ + ExecuteReport Command: Test chart email report schedule with CSV (no query context) + """ + # setup screenshot mock + screenshot_mock.return_value = SCREENSHOT_FILE + + # setup csv mock + response = Mock() + mock_open.return_value = response + mock_urlopen.return_value = response + mock_urlopen.return_value.getcode.return_value = 200 + response.read.return_value = CSV_FILE + + with freeze_time("2020-01-01T00:00:00Z"): + AsyncExecuteReportScheduleCommand( + TEST_ID, + create_report_email_chart_with_csv_no_query_context.id, + datetime.utcnow(), + ).run() + + # verify that when query context is null we request a screenshot + screenshot_mock.assert_called_once() + + @pytest.mark.usefixtures( "load_birth_names_dashboard_with_slices", "create_report_email_dashboard" )