feat: add statsd metrics for notifications (#20158)

* feat: add statsd metrics for notifications

* fix import

* fix lint

* add decorator arg for custom prefix

* add tests

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("reports.email.send")
     def send(self) -> None:
         subject = self._get_subject()
         content = self._get_content()

superset/reports/notifications/slack.py

@@ -29,6 +29,7 @@
 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 @@ 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("reports.slack.send")
     def send(self) -> None:
         files = self._get_inline_files()
         title = self._content.name

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,6 +32,27 @@
     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."""

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(

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)