feat: send data embedded in report email (#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
2021-07-28 09:43:04 -07:00 · 2021-07-28 09:43:04 -07:00 · 3adf8e85cd
parent 32d2aa0c40
commit 3adf8e85cd
10 changed files with 155 additions and 15 deletions
--- 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();
  });
--- 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>
--- 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;
--- 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):
--- a/superset/reports/api.py
+++ b/superset/reports/api.py
@ -83,6 +83,7 @@ class ReportScheduleRestApi(BaseSupersetModelRestApi):
        "active",
        "chart.id",
        "chart.slice_name",
+        "chart.viz_type",
        "context_markdown",
        "creation_method",
        "crontab",
--- 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,6 +242,14 @@ class BaseReportState:
            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
@ -247,6 +257,7 @@ class BaseReportState:
        :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 @@ class BaseReportState:
                    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 @@ class BaseReportState:
            screenshot=screenshot_data,
            description=self._report_schedule.description,
            csv=csv_data,
+            embedded_data=embedded_data,
        )

    def _send(
--- 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
--- a/superset/reports/notifications/email.py
+++ b/superset/reports/notifications/email.py
@ -21,6 +21,7 @@ from dataclasses import dataclass
 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 @@ from superset.utils.core import send_email_smtp

 logger = logging.getLogger(__name__)

+TABLE_TAGS = ["table", "th", "tr", "td", "thead", "tbody", "tfoot"]
+

@dataclass
 class EmailContent:
@ -68,14 +71,23 @@ class EmailNotification(BaseNotification):  # pylint: disable=too-few-public-met
        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 "",
--- a/tests/integration_tests/reports/api_tests.py
+++ b/tests/integration_tests/reports/api_tests.py
@ -181,6 +181,7 @@ class TestReportSchedulesApi(SupersetTestCase):
            "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,
--- 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"
 )