feat: call screenshot to store query_context (apache#15846)

* feat: call screenshot to store query_context

* Add unit test

* Move updateQueryContext to ExploreChartPanel

* Add error handling

* Fix code

* Fix logic

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;

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

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;

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

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:

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