diff --git a/superset/reports/notifications/email.py b/superset/reports/notifications/email.py index 20afeae437..3991f24b92 100644 --- a/superset/reports/notifications/email.py +++ b/superset/reports/notifications/email.py @@ -30,6 +30,7 @@ from superset.models.reports import ReportRecipientType from superset.reports.notifications.base import BaseNotification from superset.reports.notifications.exceptions import NotificationError from superset.utils.core import send_email_smtp +from superset.utils.decorators import statsd_gauge from superset.utils.urls import modify_url_query logger = logging.getLogger(__name__) @@ -149,6 +150,7 @@ class EmailNotification(BaseNotification): # pylint: disable=too-few-public-met def _get_to(self) -> str: return json.loads(self._recipient.recipient_config_json)["target"] + @statsd_gauge("reports.email.send") def send(self) -> None: subject = self._get_subject() content = self._get_content() diff --git a/superset/reports/notifications/slack.py b/superset/reports/notifications/slack.py index b833cbd53d..2a198d6645 100644 --- a/superset/reports/notifications/slack.py +++ b/superset/reports/notifications/slack.py @@ -29,6 +29,7 @@ from superset import app from superset.models.reports import ReportRecipientType from superset.reports.notifications.base import BaseNotification from superset.reports.notifications.exceptions import NotificationError +from superset.utils.decorators import statsd_gauge from superset.utils.urls import modify_url_query logger = logging.getLogger(__name__) @@ -147,6 +148,7 @@ Error: %(text)s return [] @backoff.on_exception(backoff.expo, SlackApiError, factor=10, base=2, max_tries=5) + @statsd_gauge("reports.slack.send") def send(self) -> None: files = self._get_inline_files() title = self._content.name diff --git a/superset/utils/decorators.py b/superset/utils/decorators.py index ab4ee30878..f14335f2ca 100644 --- a/superset/utils/decorators.py +++ b/superset/utils/decorators.py @@ -19,7 +19,7 @@ from __future__ import annotations import time from contextlib import contextmanager from functools import wraps -from typing import Any, Callable, Dict, Iterator, TYPE_CHECKING, Union +from typing import Any, Callable, Dict, Iterator, Optional, TYPE_CHECKING, Union from flask import current_app, Response @@ -32,6 +32,27 @@ if TYPE_CHECKING: from superset.stats_logger import BaseStatsLogger +def statsd_gauge(metric_prefix: Optional[str] = None) -> Callable[..., Any]: + def decorate(f: Callable[..., Any]) -> Callable[..., Any]: + """ + Handle sending statsd gauge metric from any method or function + """ + + def wrapped(*args: Any, **kwargs: Any) -> Any: + metric_prefix_ = metric_prefix or f.__name__ + try: + result = f(*args, **kwargs) + current_app.config["STATS_LOGGER"].gauge(f"{metric_prefix_}.ok", 1) + return result + except Exception as ex: + current_app.config["STATS_LOGGER"].gauge(f"{metric_prefix_}.error", 1) + raise ex + + return wrapped + + return decorate + + @contextmanager def stats_timing(stats_key: str, stats_logger: BaseStatsLogger) -> Iterator[float]: """Provide a transactional scope around a series of operations.""" diff --git a/tests/integration_tests/reports/commands_tests.py b/tests/integration_tests/reports/commands_tests.py index 7629bdd5b4..dd23d291fd 100644 --- a/tests/integration_tests/reports/commands_tests.py +++ b/tests/integration_tests/reports/commands_tests.py @@ -22,6 +22,7 @@ from unittest.mock import Mock, patch from uuid import uuid4 import pytest +from flask import current_app from flask_sqlalchemy import BaseQuery from freezegun import freeze_time from sqlalchemy.sql import func @@ -1026,20 +1027,23 @@ def test_email_dashboard_report_schedule( screenshot_mock.return_value = SCREENSHOT_FILE with freeze_time("2020-01-01T00:00:00Z"): - AsyncExecuteReportScheduleCommand( - TEST_ID, create_report_email_dashboard.id, datetime.utcnow() - ).run() + with patch.object(current_app.config["STATS_LOGGER"], "gauge") as statsd_mock: - notification_targets = get_target_from_report_schedule( - create_report_email_dashboard - ) - # Assert the email smtp address - assert email_mock.call_args[0][0] == notification_targets[0] - # Assert the email inline screenshot - smtp_images = email_mock.call_args[1]["images"] - assert smtp_images[list(smtp_images.keys())[0]] == SCREENSHOT_FILE - # Assert logs are correct - assert_log(ReportState.SUCCESS) + AsyncExecuteReportScheduleCommand( + TEST_ID, create_report_email_dashboard.id, datetime.utcnow() + ).run() + + notification_targets = get_target_from_report_schedule( + create_report_email_dashboard + ) + # Assert the email smtp address + assert email_mock.call_args[0][0] == notification_targets[0] + # Assert the email inline screenshot + smtp_images = email_mock.call_args[1]["images"] + assert smtp_images[list(smtp_images.keys())[0]] == SCREENSHOT_FILE + # Assert logs are correct + assert_log(ReportState.SUCCESS) + statsd_mock.assert_called_once_with("reports.email.send.ok", 1) @pytest.mark.usefixtures( @@ -1094,19 +1098,22 @@ def test_slack_chart_report_schedule( screenshot_mock.return_value = SCREENSHOT_FILE with freeze_time("2020-01-01T00:00:00Z"): - AsyncExecuteReportScheduleCommand( - TEST_ID, create_report_slack_chart.id, datetime.utcnow() - ).run() + with patch.object(current_app.config["STATS_LOGGER"], "gauge") as statsd_mock: - notification_targets = get_target_from_report_schedule( - create_report_slack_chart - ) + AsyncExecuteReportScheduleCommand( + TEST_ID, create_report_slack_chart.id, datetime.utcnow() + ).run() - assert file_upload_mock.call_args[1]["channels"] == notification_targets[0] - assert file_upload_mock.call_args[1]["file"] == SCREENSHOT_FILE + notification_targets = get_target_from_report_schedule( + create_report_slack_chart + ) - # Assert logs are correct - assert_log(ReportState.SUCCESS) + assert file_upload_mock.call_args[1]["channels"] == notification_targets[0] + assert file_upload_mock.call_args[1]["file"] == SCREENSHOT_FILE + + # Assert logs are correct + assert_log(ReportState.SUCCESS) + statsd_mock.assert_called_once_with("reports.slack.send.ok", 1) @pytest.mark.usefixtures( diff --git a/tests/integration_tests/utils/decorators_tests.py b/tests/integration_tests/utils/decorators_tests.py index 98faa8a2ba..d0ab6f9843 100644 --- a/tests/integration_tests/utils/decorators_tests.py +++ b/tests/integration_tests/utils/decorators_tests.py @@ -14,7 +14,10 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from unittest.mock import call, Mock +from unittest.mock import call, Mock, patch + +import pytest +from flask import current_app from superset.utils import decorators from tests.integration_tests.base_tests import SupersetTestCase @@ -41,3 +44,18 @@ class UtilsDecoratorsTests(SupersetTestCase): result = myfunc(1, 0, kwarg1="haha", kwarg2=2) mock.assert_has_calls([call(1, "abc"), call(1, "haha")]) self.assertEqual(result, 3) + + def test_statsd_gauge(self): + @decorators.statsd_gauge("custom.prefix") + def my_func(fail: bool, *args, **kwargs): + if fail: + raise ValueError("Error") + return "OK" + + with patch.object(current_app.config["STATS_LOGGER"], "gauge") as mock: + my_func(False, 1, 2) + mock.assert_called_once_with("custom.prefix.ok", 1) + + with pytest.raises(ValueError): + my_func(True, 1, 2) + mock.assert_called_once_with("custom.prefix.error", 1)