This is an automated email from the ASF dual-hosted git repository. ephraimanierobi pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push: new e4ca68818e Adapt Notifier for sla_miss_callback (#31887) e4ca68818e is described below commit e4ca68818eec0f29ef04a1a5bfec3241ea03bf8c Author: Utkarsh Sharma <utkarshar...@gmail.com> AuthorDate: Fri Jun 23 22:31:18 2023 +0530 Adapt Notifier for sla_miss_callback (#31887) * Fix Notifier issue with sla_miss_callback * Add testcase to ensure the old pattern is working * Update document for existing notifier * Made notifier backward compatible with sla_miss_callback * Remove unwanted imports * Fix testcases * Fix static check * Fix PR comments * Removed unwanted code * Remove uneanted change * Updated var name * Remove kwargs * Fix documentation * Fix tests --- airflow/notifications/basenotifier.py | 18 +++++++++++-- .../pagerduty_notifier_howto_guide.rst | 5 ++++ .../notifications/slack_notifier_howto_guide.rst | 5 ++++ .../notifications/smtp_notifier_howto_guide.rst | 5 ++++ tests/notifications/test_basenotifier.py | 30 ++++++++++++++++++++++ .../apprise/notifications/test_apprise.py | 4 +-- .../pagerduty/notifications/test_pagerduty.py | 4 +-- tests/providers/slack/notifications/test_slack.py | 4 +-- tests/providers/smtp/notifications/test_smtp.py | 4 +-- 9 files changed, 69 insertions(+), 10 deletions(-) diff --git a/airflow/notifications/basenotifier.py b/airflow/notifications/basenotifier.py index 5f922812d4..7ef0603be1 100644 --- a/airflow/notifications/basenotifier.py +++ b/airflow/notifications/basenotifier.py @@ -79,13 +79,27 @@ class BaseNotifier(Templater): """ ... - def __call__(self, context: Context) -> None: + def __call__(self, *args) -> None: """ Send a notification. :param context: The airflow context """ - context = self._update_context(context) + # Currently, there are two ways a callback is invoked + # 1. callback(context) - for on_*_callbacks + # 2. callback(dag, task_list, blocking_task_list, slas, blocking_tis) - for sla_miss_callback + # we have to distinguish between the two calls so that we can prepare the correct context, + if len(args) == 1: + context = args[0] + else: + context = { + "dag": args[0], + "task_list": args[1], + "blocking_task_list": args[2], + "slas": args[3], + "blocking_tis": args[4], + } + self._update_context(context) self.render_template_fields(context) try: self.notify(context) diff --git a/docs/apache-airflow-providers-pagerduty/notifications/pagerduty_notifier_howto_guide.rst b/docs/apache-airflow-providers-pagerduty/notifications/pagerduty_notifier_howto_guide.rst index 2a2388f692..4d6fd4c929 100644 --- a/docs/apache-airflow-providers-pagerduty/notifications/pagerduty_notifier_howto_guide.rst +++ b/docs/apache-airflow-providers-pagerduty/notifications/pagerduty_notifier_howto_guide.rst @@ -23,6 +23,11 @@ Introduction The Pagerduty notifier (:class:`airflow.providers.pagerduty.notifications.pagerduty.PagerdutyNotifier`) allows users to send messages to Pagerduty using the various ``on_*_callbacks`` at both the DAG level and Task level. +You can also use a notifier with ``sla_miss_callback``. + +.. note:: + When notifiers are used with `sla_miss_callback` the context will contain only values passed to the callback, refer :ref:`sla_miss_callback<concepts:sla_miss_callback>`. + Example Code: ------------- diff --git a/docs/apache-airflow-providers-slack/notifications/slack_notifier_howto_guide.rst b/docs/apache-airflow-providers-slack/notifications/slack_notifier_howto_guide.rst index c6b087398c..d967779cee 100644 --- a/docs/apache-airflow-providers-slack/notifications/slack_notifier_howto_guide.rst +++ b/docs/apache-airflow-providers-slack/notifications/slack_notifier_howto_guide.rst @@ -23,6 +23,11 @@ Introduction Slack notifier (:class:`airflow.providers.slack.notifications.slack.SlackNotifier`) allows users to send messages to a slack channel using the various ``on_*_callbacks`` at both the DAG level and Task level +You can also use a notifier with ``sla_miss_callback``. + +.. note:: + When notifiers are used with `sla_miss_callback` the context will contain only values passed to the callback, refer :ref:`sla_miss_callback<concepts:sla_miss_callback>`. + Example Code: ------------- diff --git a/docs/apache-airflow-providers-smtp/notifications/smtp_notifier_howto_guide.rst b/docs/apache-airflow-providers-smtp/notifications/smtp_notifier_howto_guide.rst index 9d9d0ac5ee..99895aef3c 100644 --- a/docs/apache-airflow-providers-smtp/notifications/smtp_notifier_howto_guide.rst +++ b/docs/apache-airflow-providers-smtp/notifications/smtp_notifier_howto_guide.rst @@ -23,6 +23,11 @@ Introduction The SMTP notifier (:class:`airflow.providers.smtp.notifications.smtp.SmtpNotifier`) allows users to send messages to SMTP servers using the various ``on_*_callbacks`` at both the DAG level and Task level. +You can also use a notifier with ``sla_miss_callback``. + +.. note:: + When notifiers are used with `sla_miss_callback` the context will contain only values passed to the callback, refer :ref:`sla_miss_callback<concepts:sla_miss_callback>`. + Example Code: ------------- diff --git a/tests/notifications/test_basenotifier.py b/tests/notifications/test_basenotifier.py index 574d577f75..cb768890d7 100644 --- a/tests/notifications/test_basenotifier.py +++ b/tests/notifications/test_basenotifier.py @@ -17,6 +17,8 @@ from __future__ import annotations +from unittest.mock import MagicMock + import jinja2 import pytest @@ -64,3 +66,31 @@ class TestBaseNotifier: context: Context = {"dag": dag} notifier.render_template_fields(context) assert notifier.message == "Hello test_render_message_with_template_works" + + def test_notifier_call_with_passed_context(self, dag_maker, caplog): + with dag_maker("test_render_message_with_template_works") as dag: + EmptyOperator(task_id="test_id") + notifier = MockNotifier(message="Hello {{ dag.dag_id }}") + notifier.notify = MagicMock() + context: Context = {"dag": dag} + notifier(context) + notifier.notify.assert_called_once_with({"dag": dag, "message": "Hello {{ dag.dag_id }}"}) + assert notifier.message == "Hello test_render_message_with_template_works" + + def test_notifier_call_with_prepared_context(self, dag_maker, caplog): + with dag_maker("test_render_message_with_template_works"): + EmptyOperator(task_id="test_id") + notifier = MockNotifier(message="task: {{ task_list[0] }}") + notifier.notify = MagicMock() + notifier(None, ["some_task"], None, None, None) + notifier.notify.assert_called_once_with( + { + "dag": None, + "task_list": ["some_task"], + "blocking_task_list": None, + "slas": None, + "blocking_tis": None, + "message": "task: {{ task_list[0] }}", + } + ) + assert notifier.message == "task: some_task" diff --git a/tests/providers/apprise/notifications/test_apprise.py b/tests/providers/apprise/notifications/test_apprise.py index c7dc837419..c81fcbd332 100644 --- a/tests/providers/apprise/notifications/test_apprise.py +++ b/tests/providers/apprise/notifications/test_apprise.py @@ -34,7 +34,7 @@ class TestAppriseNotifier: with dag_maker("test_notifier") as dag: EmptyOperator(task_id="task1") notifier = send_apprise_notification(body="DISK at 99%", notify_type=NotifyType.FAILURE) - notifier(context={"dag": dag}) + notifier({"dag": dag}) mock_apprise_hook.return_value.notify.assert_called_once_with( body="DISK at 99%", notify_type=NotifyType.FAILURE, @@ -51,7 +51,7 @@ class TestAppriseNotifier: with dag_maker("test_notifier") as dag: EmptyOperator(task_id="task1") notifier = AppriseNotifier(body="DISK at 99%", notify_type=NotifyType.FAILURE) - notifier(context={"dag": dag}) + notifier({"dag": dag}) mock_apprise_hook.return_value.notify.assert_called_once_with( body="DISK at 99%", notify_type=NotifyType.FAILURE, diff --git a/tests/providers/pagerduty/notifications/test_pagerduty.py b/tests/providers/pagerduty/notifications/test_pagerduty.py index 544c4d217d..369760c3bf 100644 --- a/tests/providers/pagerduty/notifications/test_pagerduty.py +++ b/tests/providers/pagerduty/notifications/test_pagerduty.py @@ -35,7 +35,7 @@ class TestPagerdutyNotifier: with dag_maker("test_notifier") as dag: EmptyOperator(task_id="task1") notifier = send_pagerduty_notification(summary="DISK at 99%", severity="critical", action="trigger") - notifier(context={"dag": dag}) + notifier({"dag": dag}) mock_pagerduty_event_hook.return_value.create_event.assert_called_once_with( summary="DISK at 99%", severity="critical", @@ -55,7 +55,7 @@ class TestPagerdutyNotifier: with dag_maker("test_notifier") as dag: EmptyOperator(task_id="task1") notifier = PagerdutyNotifier(summary="DISK at 99%", severity="critical", action="trigger") - notifier(context={"dag": dag}) + notifier({"dag": dag}) mock_pagerduty_event_hook.return_value.create_event.assert_called_once_with( summary="DISK at 99%", severity="critical", diff --git a/tests/providers/slack/notifications/test_slack.py b/tests/providers/slack/notifications/test_slack.py index f8c22fe0d3..295a01d092 100644 --- a/tests/providers/slack/notifications/test_slack.py +++ b/tests/providers/slack/notifications/test_slack.py @@ -30,7 +30,7 @@ class TestSlackNotifier: EmptyOperator(task_id="task1") notifier = send_slack_notification(text="test") - notifier(context={"dag": dag}) + notifier({"dag": dag}) mock_slack_hook.return_value.call.assert_called_once_with( "chat.postMessage", json={ @@ -50,7 +50,7 @@ class TestSlackNotifier: EmptyOperator(task_id="task1") notifier = SlackNotifier(text="test") - notifier(context={"dag": dag}) + notifier({"dag": dag}) mock_slack_hook.return_value.call.assert_called_once_with( "chat.postMessage", json={ diff --git a/tests/providers/smtp/notifications/test_smtp.py b/tests/providers/smtp/notifications/test_smtp.py index 44d2ce6b9d..11578022d8 100644 --- a/tests/providers/smtp/notifications/test_smtp.py +++ b/tests/providers/smtp/notifications/test_smtp.py @@ -40,7 +40,7 @@ class TestPagerdutyNotifier: subject="subject", html_content="body", ) - notifier(context={"dag": dag}) + notifier({"dag": dag}) mock_smtphook_hook.return_value.__enter__().send_email_smtp.assert_called_once_with( from_email="test_sen...@test.com", to="test_reci...@test.com", @@ -65,7 +65,7 @@ class TestPagerdutyNotifier: subject="subject", html_content="body", ) - notifier(context={"dag": dag}) + notifier({"dag": dag}) mock_smtphook_hook.return_value.__enter__().send_email_smtp.assert_called_once_with( from_email="test_sen...@test.com", to="test_reci...@test.com",