feat: send data embedded in report email (apache#15805)

* feat: send data embedded in report email * Prettify table * Change post-processing to use new endpoint * Show text option only for text viz * Show TEXT option only to text-based vizs * Fix test * Add email test * Add unit test
CybercentreCanada · Dec 17, 2021 · 8b31066 · 8b31066
1 parent ff0d88c
commit 8b31066
Show file tree

Hide file tree

Showing 10 changed files with 155 additions and 15 deletions.
diff --git a/superset-frontend/src/views/CRUD/alert/AlertReportModal.test.jsx b/superset-frontend/src/views/CRUD/alert/AlertReportModal.test.jsx
@@ -36,7 +36,7 @@ const mockData = {
   id: 1,
   name: 'test report',
   description: 'test report description',
-  chart: { id: 1, slice_name: 'test chart' },
+  chart: { id: 1, slice_name: 'test chart', viz_type: 'table' },
   database: { id: 1, database_name: 'test database' },
   sql: 'SELECT NaN',
 };
@@ -76,7 +76,7 @@ fetchMock.get(dashboardEndpoint, {
 });
 
 fetchMock.get(chartEndpoint, {
-  result: [],
+  result: [{ text: 'table chart', value: 1 }],
 });
 
 async function mountAndWait(props = mockedProps) {
@@ -260,6 +260,22 @@ describe('AlertReportModal', () => {
     expect(wrapper.find(Radio)).toHaveLength(2);
   });
 
+  it('renders text option for text-based charts', async () => {
+    const props = {
+      ...mockedProps,
+      alert: mockData,
+    };
+    const textWrapper = await mountAndWait(props);
+
+    const chartOption = textWrapper.find('input[value="chart"]');
+    act(() => {
+      chartOption.props().onChange({ target: { value: 'chart' } });
+    });
+    await waitForComponentToPaint(textWrapper);
+
+    expect(textWrapper.find('input[value="TEXT"]')).toExist();
+  });
+
   it('renders input element for working timeout', () => {
     expect(wrapper.find('input[name="working_timeout"]')).toExist();
   });

diff --git a/superset-frontend/src/views/CRUD/alert/AlertReportModal.tsx b/superset-frontend/src/views/CRUD/alert/AlertReportModal.tsx
@@ -51,6 +51,12 @@ import {
 
 const SELECT_PAGE_SIZE = 2000; // temporary fix for paginated query
 const TIMEOUT_MIN = 1;
+const TEXT_BASED_VISUALIZATION_TYPES = [
+  'pivot_table',
+  'pivot_table_v2',
+  'table',
+  'paired_ttest',
+];
 
 type SelectValue = {
   value: string;
@@ -416,6 +422,9 @@ const AlertReportModal: FunctionComponent<AlertReportModalProps> = ({
   const [dashboardOptions, setDashboardOptions] = useState<MetaObject[]>([]);
   const [chartOptions, setChartOptions] = useState<MetaObject[]>([]);
 
+  // Chart metadata
+  const [chartVizType, setChartVizType] = useState<string>('');
+
   const isEditMode = alert !== null;
   const formatOptionEnabled =
     contentType === 'chart' &&
@@ -718,6 +727,11 @@ const AlertReportModal: FunctionComponent<AlertReportModalProps> = ({
     return result;
   };
 
+  const getChartVisualizationType = (chart: SelectValue) =>
+    SupersetClient.get({
+      endpoint: `/api/v1/chart/${chart.value}`,
+    }).then(response => setChartVizType(response.json.result.viz_type));
+
   // Updating alert/report state
   const updateAlertState = (name: string, value: any) => {
     setCurrentAlert(currentAlertData => ({
@@ -770,6 +784,7 @@ const AlertReportModal: FunctionComponent<AlertReportModalProps> = ({
   };
 
   const onChartChange = (chart: SelectValue) => {
+    getChartVisualizationType(chart);
     updateAlertState('chart', chart || undefined);
     updateAlertState('dashboard', null);
   };
@@ -920,6 +935,10 @@ const AlertReportModal: FunctionComponent<AlertReportModalProps> = ({
 
       setConditionNotNull(resource.validator_type === 'not null');
 
+      if (resource.chart) {
+        setChartVizType((resource.chart as ChartObject).viz_type);
+      }
+
       setCurrentAlert({
         ...resource,
         chart: resource.chart
@@ -1251,17 +1270,6 @@ const AlertReportModal: FunctionComponent<AlertReportModalProps> = ({
               <StyledRadio value="dashboard">{t('Dashboard')}</StyledRadio>
               <StyledRadio value="chart">{t('Chart')}</StyledRadio>
             </Radio.Group>
-            {formatOptionEnabled && (
-              <div className="inline-container">
-                <StyledRadioGroup
-                  onChange={onFormatChange}
-                  value={reportFormat}
-                >
-                  <StyledRadio value="PNG">{t('Send as PNG')}</StyledRadio>
-                  <StyledRadio value="CSV">{t('Send as CSV')}</StyledRadio>
-                </StyledRadioGroup>
-              </div>
-            )}
             <AsyncSelect
               className={
                 contentType === 'chart'
@@ -1302,6 +1310,20 @@ const AlertReportModal: FunctionComponent<AlertReportModalProps> = ({
               cacheOptions
               onChange={onDashboardChange}
             />
+            {formatOptionEnabled && (
+              <div className="inline-container">
+                <StyledRadioGroup
+                  onChange={onFormatChange}
+                  value={reportFormat}
+                >
+                  <StyledRadio value="PNG">{t('Send as PNG')}</StyledRadio>
+                  <StyledRadio value="CSV">{t('Send as CSV')}</StyledRadio>
+                  {TEXT_BASED_VISUALIZATION_TYPES.includes(chartVizType) && (
+                    <StyledRadio value="TEXT">{t('Send as text')}</StyledRadio>
+                  )}
+                </StyledRadioGroup>
+              </div>
+            )}
             <StyledSectionTitle>
               <h4>{t('Notification method')}</h4>
               <span className="required">*</span>

diff --git a/superset-frontend/src/views/CRUD/alert/types.ts b/superset-frontend/src/views/CRUD/alert/types.ts
@@ -27,6 +27,7 @@ type user = {
 export type ChartObject = {
   id: number;
   slice_name: string;
+  viz_type: string;
 };
 
 export type DashboardObject = {
@@ -74,7 +75,7 @@ export type AlertObject = {
   owners?: Array<Owner | MetaObject>;
   sql?: string;
   recipients?: Array<Recipient>;
-  report_format?: 'PNG' | 'CSV';
+  report_format?: 'PNG' | 'CSV' | 'TEXT';
   type?: string;
   validator_config_json?: {
     op?: Operator;

diff --git a/superset/models/reports.py b/superset/models/reports.py
@@ -72,6 +72,7 @@ class ReportState(str, enum.Enum):
 class ReportDataFormat(str, enum.Enum):
     VISUALIZATION = "PNG"
     DATA = "CSV"
+    TEXT = "TEXT"
 
 
 class ReportCreationMethodType(str, enum.Enum):

diff --git a/superset/reports/api.py b/superset/reports/api.py
@@ -83,6 +83,7 @@ def ensure_alert_reports_enabled(self) -> Optional[Response]:
         "active",
         "chart.id",
         "chart.slice_name",
+        "chart.viz_type",
         "context_markdown",
         "creation_method",
         "crontab",

diff --git a/superset/reports/commands/execute.py b/superset/reports/commands/execute.py
@@ -17,9 +17,11 @@
 import json
 import logging
 from datetime import datetime, timedelta
+from io import BytesIO
 from typing import Any, List, Optional
 from uuid import UUID
 
+import pandas as pd
 from celery.exceptions import SoftTimeLimitExceeded
 from flask_appbuilder.security.sqla.models import User
 from sqlalchemy.orm import Session
@@ -240,13 +242,22 @@ def _get_csv_data(self) -> bytes:
             raise ReportScheduleCsvFailedError()
         return csv_data
 
+    def _get_embedded_data(self) -> str:
+        """
+        Return data as an HTML table, to embed in the email.
+        """
+        buf = BytesIO(self._get_csv_data())
+        df = pd.read_csv(buf)
+        return df.to_html(na_rep="", index=False)
+
     def _get_notification_content(self) -> NotificationContent:
         """
         Gets a notification content, this is composed by a title and a screenshot
 
         :raises: ReportScheduleScreenshotFailedError
         """
         csv_data = None
+        embedded_data = None
         error_text = None
         screenshot_data = None
         url = self._get_url(user_friendly=True)
@@ -270,6 +281,12 @@ def _get_notification_content(self) -> NotificationContent:
                     name=self._report_schedule.name, text=error_text
                 )
 
+        if (
+            self._report_schedule.chart
+            and self._report_schedule.report_format == ReportDataFormat.TEXT
+        ):
+            embedded_data = self._get_embedded_data()
+
         if self._report_schedule.chart:
             name = (
                 f"{self._report_schedule.name}: "
@@ -286,6 +303,7 @@ def _get_notification_content(self) -> NotificationContent:
             screenshot=screenshot_data,
             description=self._report_schedule.description,
             csv=csv_data,
+            embedded_data=embedded_data,
         )
 
     def _send(

diff --git a/superset/reports/notifications/base.py b/superset/reports/notifications/base.py
@@ -29,6 +29,7 @@ class NotificationContent:
     text: Optional[str] = None
     description: Optional[str] = ""
     url: Optional[str] = None  # url to chart/dashboard for this screenshot
+    embedded_data: Optional[str] = ""
 
 
 class BaseNotification:  # pylint: disable=too-few-public-methods

diff --git a/superset/reports/notifications/email.py b/superset/reports/notifications/email.py
@@ -21,6 +21,7 @@
 from email.utils import make_msgid, parseaddr
 from typing import Any, Dict, Optional
 
+import bleach
 from flask_babel import gettext as __
 
 from superset import app
@@ -31,6 +32,8 @@
 
 logger = logging.getLogger(__name__)
 
+TABLE_TAGS = ["table", "th", "tr", "td", "thead", "tbody", "tfoot"]
+
 
 @dataclass
 class EmailContent:
@@ -68,14 +71,23 @@ def _get_content(self) -> EmailContent:
         csv_data = None
         domain = self._get_smtp_domain()
         msgid = make_msgid(domain)[1:-1]
+
+        # Strip any malicious HTML from the description
+        description = bleach.clean(self._content.description or "")
+
+        # Strip malicious HTML from embedded data, allowing table elements
+        embedded_data = bleach.clean(self._content.embedded_data or "", tags=TABLE_TAGS)
+
         body = __(
             """
             <p>%(description)s</p>
             <b><a href="%(url)s">Explore in Superset</a></b><p></p>
+            %(embedded_data)s
             %(img_tag)s
             """,
-            description=self._content.description or "",
+            description=description,
             url=self._content.url,
+            embedded_data=embedded_data,
             img_tag='<img width="1000px" src="cid:{}">'.format(msgid)
             if self._content.screenshot
             else "",

diff --git a/tests/integration_tests/reports/api_tests.py b/tests/integration_tests/reports/api_tests.py
@@ -181,6 +181,7 @@ def test_get_report_schedule(self):
             "chart": {
                 "id": report_schedule.chart.id,
                 "slice_name": report_schedule.chart.slice_name,
+                "viz_type": report_schedule.chart.viz_type,
             },
             "context_markdown": report_schedule.context_markdown,
             "crontab": report_schedule.crontab,

diff --git a/tests/integration_tests/reports/commands_tests.py b/tests/integration_tests/reports/commands_tests.py
@@ -230,6 +230,20 @@ def create_report_email_chart_with_csv():
         cleanup_report_schedule(report_schedule)
 
 
+@pytest.fixture()
+def create_report_email_chart_with_text():
+    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,
+            report_format=ReportDataFormat.TEXT,
+        )
+        yield report_schedule
+        cleanup_report_schedule(report_schedule)
+
+
 @pytest.fixture()
 def create_report_email_chart_with_csv_no_query_context():
     with app.app_context():
@@ -721,6 +735,59 @@ def test_email_chart_report_schedule_with_csv_no_query_context(
         screenshot_mock.assert_called_once()
 
 
+@pytest.mark.usefixtures(
+    "load_birth_names_dashboard_with_slices", "create_report_email_chart_with_text"
+)
+@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")
+def test_email_chart_report_schedule_with_text(
+    csv_mock, email_mock, mock_open, mock_urlopen, create_report_email_chart_with_text,
+):
+    """
+    ExecuteReport Command: Test chart email report schedule with CSV
+    """
+    # 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_text.id, datetime.utcnow()
+        ).run()
+
+        # assert that the data is embedded correctly
+        table_html = """<table>
+  <thead>
+    <tr>
+      <th>t1</th>
+      <th>t2</th>
+      <th>t3__sum</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td>c11</td>
+      <td>c12</td>
+      <td>c13</td>
+    </tr>
+    <tr>
+      <td>c21</td>
+      <td>c22</td>
+      <td>c23</td>
+    </tr>
+  </tbody>
+</table>"""
+        assert table_html in email_mock.call_args[0][2]
+
+        # Assert logs are correct
+        assert_log(ReportState.SUCCESS)
+
+
 @pytest.mark.usefixtures(
     "load_birth_names_dashboard_with_slices", "create_report_email_dashboard"
 )