From be14014e5be355d88007b999ae60ddd5ec6b4463 Mon Sep 17 00:00:00 2001 From: dpgaspar Date: Mon, 23 May 2022 10:13:15 +0100 Subject: [PATCH 1/5] feat: add statsd metrics for notifications --- superset/reports/notifications/email.py | 2 ++ superset/reports/notifications/slack.py | 2 ++ superset/utils/decorators.py | 17 +++++++++++++++++ 3 files changed, 21 insertions(+) diff --git a/superset/reports/notifications/email.py b/superset/reports/notifications/email.py index 20afeae437b00..31873ddefe32d 100644 --- a/superset/reports/notifications/email.py +++ b/superset/reports/notifications/email.py @@ -30,6 +30,7 @@ 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 @@ def _get_subject(self) -> str: def _get_to(self) -> str: return json.loads(self._recipient.recipient_config_json)["target"] + @statsd_gauge 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 b833cbd53ddf3..46f9eded6c816 100644 --- a/superset/reports/notifications/slack.py +++ b/superset/reports/notifications/slack.py @@ -24,6 +24,7 @@ from flask_babel import gettext as __ from slack import WebClient from slack.errors import SlackApiError, SlackClientError +from utils.decorators import statsd_gauge from superset import app from superset.models.reports import ReportRecipientType @@ -147,6 +148,7 @@ def _get_inline_files(self) -> Sequence[Union[str, IOBase, bytes]]: return [] @backoff.on_exception(backoff.expo, SlackApiError, factor=10, base=2, max_tries=5) + @statsd_gauge 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 ab4ee308787c8..f93a30c089c93 100644 --- a/superset/utils/decorators.py +++ b/superset/utils/decorators.py @@ -32,6 +32,23 @@ from superset.stats_logger import BaseStatsLogger +def statsd_gauge(f: Callable[..., Any]) -> Callable[..., Any]: + """ + Handle sending statsd gauge metric from any method or function + """ + + def wraps(*args: Any, **kwargs: Any) -> Any: + try: + result = f(*args, **kwargs) + current_app.config["STATS_LOGGER"].gauge(f"{f.__name__}.ok", 1) + return result + except Exception as ex: + current_app.config["STATS_LOGGER"].gauge(f"{f.__name__}.error", 1) + raise ex + + return wraps + + @contextmanager def stats_timing(stats_key: str, stats_logger: BaseStatsLogger) -> Iterator[float]: """Provide a transactional scope around a series of operations.""" From 29250b559235f9d4388ad369802fc8b1544570cd Mon Sep 17 00:00:00 2001 From: dpgaspar Date: Mon, 23 May 2022 10:29:29 +0100 Subject: [PATCH 2/5] fix import --- superset/reports/notifications/slack.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/superset/reports/notifications/slack.py b/superset/reports/notifications/slack.py index 46f9eded6c816..4e346e8d16c16 100644 --- a/superset/reports/notifications/slack.py +++ b/superset/reports/notifications/slack.py @@ -24,12 +24,12 @@ from flask_babel import gettext as __ from slack import WebClient from slack.errors import SlackApiError, SlackClientError -from utils.decorators import statsd_gauge 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__) From 2b3ca8faa66c25c180e2a60242cbe1336fa6a38c Mon Sep 17 00:00:00 2001 From: dpgaspar Date: Mon, 23 May 2022 10:53:04 +0100 Subject: [PATCH 3/5] fix lint --- superset/utils/decorators.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/superset/utils/decorators.py b/superset/utils/decorators.py index f93a30c089c93..21f855418e0aa 100644 --- a/superset/utils/decorators.py +++ b/superset/utils/decorators.py @@ -37,7 +37,7 @@ def statsd_gauge(f: Callable[..., Any]) -> Callable[..., Any]: Handle sending statsd gauge metric from any method or function """ - def wraps(*args: Any, **kwargs: Any) -> Any: + def wrapped(*args: Any, **kwargs: Any) -> Any: try: result = f(*args, **kwargs) current_app.config["STATS_LOGGER"].gauge(f"{f.__name__}.ok", 1) @@ -46,7 +46,7 @@ def wraps(*args: Any, **kwargs: Any) -> Any: current_app.config["STATS_LOGGER"].gauge(f"{f.__name__}.error", 1) raise ex - return wraps + return wrapped @contextmanager From 6d904226dfa1300a0066008100cf39366c09949c Mon Sep 17 00:00:00 2001 From: dpgaspar Date: Mon, 23 May 2022 14:05:52 +0100 Subject: [PATCH 4/5] add decorator arg for custom prefix --- superset/reports/notifications/email.py | 2 +- superset/reports/notifications/slack.py | 2 +- superset/utils/decorators.py | 32 ++++++++++++++----------- 3 files changed, 20 insertions(+), 16 deletions(-) diff --git a/superset/reports/notifications/email.py b/superset/reports/notifications/email.py index 31873ddefe32d..0b5ebf5a64210 100644 --- a/superset/reports/notifications/email.py +++ b/superset/reports/notifications/email.py @@ -150,7 +150,7 @@ def _get_subject(self) -> str: def _get_to(self) -> str: return json.loads(self._recipient.recipient_config_json)["target"] - @statsd_gauge + @statsd_gauge("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 4e346e8d16c16..24e09e400b189 100644 --- a/superset/reports/notifications/slack.py +++ b/superset/reports/notifications/slack.py @@ -148,7 +148,7 @@ def _get_inline_files(self) -> Sequence[Union[str, IOBase, bytes]]: return [] @backoff.on_exception(backoff.expo, SlackApiError, factor=10, base=2, max_tries=5) - @statsd_gauge + @statsd_gauge("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 21f855418e0aa..f14335f2cada5 100644 --- a/superset/utils/decorators.py +++ b/superset/utils/decorators.py @@ -19,7 +19,7 @@ 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,21 +32,25 @@ from superset.stats_logger import BaseStatsLogger -def statsd_gauge(f: Callable[..., Any]) -> Callable[..., Any]: - """ - Handle sending statsd gauge metric from any method or function - """ +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: - try: - result = f(*args, **kwargs) - current_app.config["STATS_LOGGER"].gauge(f"{f.__name__}.ok", 1) - return result - except Exception as ex: - current_app.config["STATS_LOGGER"].gauge(f"{f.__name__}.error", 1) - raise ex + 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 wrapped + return decorate @contextmanager From ad9257bd72f819af3a84d247c533e0ad4ea5b2e1 Mon Sep 17 00:00:00 2001 From: dpgaspar Date: Thu, 26 May 2022 11:51:11 +0100 Subject: [PATCH 5/5] add tests --- superset/reports/notifications/email.py | 2 +- superset/reports/notifications/slack.py | 2 +- .../reports/commands_tests.py | 53 +++++++++++-------- .../utils/decorators_tests.py | 20 ++++++- 4 files changed, 51 insertions(+), 26 deletions(-) diff --git a/superset/reports/notifications/email.py b/superset/reports/notifications/email.py index 0b5ebf5a64210..3991f24b9264d 100644 --- a/superset/reports/notifications/email.py +++ b/superset/reports/notifications/email.py @@ -150,7 +150,7 @@ def _get_subject(self) -> str: def _get_to(self) -> str: return json.loads(self._recipient.recipient_config_json)["target"] - @statsd_gauge("email.send") + @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 24e09e400b189..2a198d66453c2 100644 --- a/superset/reports/notifications/slack.py +++ b/superset/reports/notifications/slack.py @@ -148,7 +148,7 @@ def _get_inline_files(self) -> Sequence[Union[str, IOBase, bytes]]: return [] @backoff.on_exception(backoff.expo, SlackApiError, factor=10, base=2, max_tries=5) - @statsd_gauge("slack.send") + @statsd_gauge("reports.slack.send") def send(self) -> None: files = self._get_inline_files() title = self._content.name diff --git a/tests/integration_tests/reports/commands_tests.py b/tests/integration_tests/reports/commands_tests.py index 7629bdd5b4583..dd23d291fd69a 100644 --- a/tests/integration_tests/reports/commands_tests.py +++ b/tests/integration_tests/reports/commands_tests.py @@ -22,6 +22,7 @@ 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 98faa8a2ba6a4..d0ab6f98434b3 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 @@ def myfunc(arg1: int, arg2: int, kwarg1: str = "abc", kwarg2: int = 2): 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)