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 <[email protected]>
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="[email protected]",
to="[email protected]",
@@ -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="[email protected]",
to="[email protected]",
