diff --git a/superset/config.py b/superset/config.py index 94abd64cd..2b72e0177 100644 --- a/superset/config.py +++ b/superset/config.py @@ -28,7 +28,7 @@ import os import sys from collections import OrderedDict from datetime import date -from typing import Any, Callable, Dict, List, Optional, Type, TYPE_CHECKING +from typing import Any, Callable, Dict, List, Optional, Type, TYPE_CHECKING, Union from cachelib.base import BaseCache from celery.schedules import crontab @@ -912,8 +912,8 @@ ALERT_REPORTS_CRON_WINDOW_SIZE = 59 # A custom prefix to use on all Alerts & Reports emails EMAIL_REPORTS_SUBJECT_PREFIX = "[Report] " -# Slack API token for the superset reports -SLACK_API_TOKEN = None +# Slack API token for the superset reports, either string or callable +SLACK_API_TOKEN: Optional[Union[Callable[[], str], str]] = None SLACK_PROXY = None # If enabled, certain features are run in debug mode diff --git a/superset/reports/notifications/slack.py b/superset/reports/notifications/slack.py index 84f858c1e..9732aad16 100644 --- a/superset/reports/notifications/slack.py +++ b/superset/reports/notifications/slack.py @@ -79,9 +79,10 @@ class SlackNotification(BaseNotification): # pylint: disable=too-few-public-met channel = self._get_channel() body = self._get_body() try: - client = WebClient( - token=app.config["SLACK_API_TOKEN"], proxy=app.config["SLACK_PROXY"] - ) + token = app.config["SLACK_API_TOKEN"] + if callable(token): + token = token() + client = WebClient(token=token, proxy=app.config["SLACK_PROXY"]) # files_upload returns SlackResponse as we run it in sync mode. if file: client.files_upload( diff --git a/superset/tasks/slack_util.py b/superset/tasks/slack_util.py index 9cb8c9179..5440698e7 100644 --- a/superset/tasks/slack_util.py +++ b/superset/tasks/slack_util.py @@ -39,7 +39,10 @@ def deliver_slack_msg( file: Optional[Union[str, IOBase, bytes]], ) -> None: config = current_app.config - client = WebClient(token=config["SLACK_API_TOKEN"], proxy=config["SLACK_PROXY"]) + token = config["SLACK_API_TOKEN"] + if callable(token): + token = token() + client = WebClient(token=token, proxy=config["SLACK_PROXY"]) # files_upload returns SlackResponse as we run it in sync mode. if file: response = cast( diff --git a/tests/reports/commands_tests.py b/tests/reports/commands_tests.py index 48c3ee230..9234adb14 100644 --- a/tests/reports/commands_tests.py +++ b/tests/reports/commands_tests.py @@ -17,7 +17,7 @@ import json from datetime import datetime, timedelta from typing import List, Optional -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest from contextlib2 import contextmanager @@ -770,6 +770,32 @@ def test_slack_chart_alert(screenshot_mock, email_mock, create_alert_email_chart assert_log(ReportState.SUCCESS) +@pytest.mark.usefixtures( + "load_birth_names_dashboard_with_slices", "create_report_slack_chart" +) +@patch("superset.reports.notifications.slack.WebClient") +@patch("superset.utils.screenshots.ChartScreenshot.get_screenshot") +def test_slack_token_callable_chart_report( + screenshot_mock, slack_client_mock_class, create_report_slack_chart +): + """ + ExecuteReport Command: Test chart slack alert (slack token callable) + """ + slack_client_mock_class.return_value = Mock() + app.config["SLACK_API_TOKEN"] = Mock(return_value="cool_code") + # setup screenshot mock + screenshot = read_fixture("sample.png") + screenshot_mock.return_value = screenshot + + with freeze_time("2020-01-01T00:00:00Z"): + AsyncExecuteReportScheduleCommand( + create_report_slack_chart.id, datetime.utcnow() + ).run() + app.config["SLACK_API_TOKEN"].assert_called_once() + assert slack_client_mock_class.called_with(token="cool_code", proxy="") + assert_log(ReportState.SUCCESS) + + @pytest.mark.usefixtures("create_no_alert_email_chart") def test_email_chart_no_alert(create_no_alert_email_chart): """