如何将参数传递给 Airflow on_success_callback 和 on_failure_callback

Question

我已经使用 on_success_callback 和 on_failure_callback 实现了关于成功和失败的电子邮件警报。

根据气流文档，

上下文字典作为单个参数传递给此函数。

如何将另一个参数传递给这些回调方法？

这是我的代码

from airflow.utils.email import send_email_smtp

def task_success_alert(context):
    subject = "[Airflow] DAG {0} - Task {1}: Success".format(
        context['task_instance_key_str'].split('__')[0], 
        context['task_instance_key_str'].split('__')[1]
        )
    html_content = """
    DAG: {0}<br>
    Task: {1}<br>
    Succeeded on: {2}
    """.format(
        context['task_instance_key_str'].split('__')[0], 
        context['task_instance_key_str'].split('__')[1], 
        datetime.now()
        )
    send_email_smtp(dag_vars["dev_mailing_list"], subject, html_content)

def task_failure_alert(context):
    subject = "[Airflow] DAG {0} - Task {1}: Failed".format(
        context['task_instance_key_str'].split('__')[0], 
        context['task_instance_key_str'].split('__')[1]
        )
    html_content = """
    DAG: {0}<br>
    Task: {1}<br>
    Failed on: {2}
    """.format(
        context['task_instance_key_str'].split('__')[0], 
        context['task_instance_key_str'].split('__')[1], 
        datetime.now()
        )
    send_email_smtp(dag_vars["dev_mailing_list"], subject, html_content)

default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2019, 6, 13),
    'on_success_callback': task_success_alert,
    'on_failure_callback': task_failure_alert
}

我打算将回调移动到另一个包并将电子邮件地址作为参数传递。

Answer 1

您可以在 dag 中定义一个函数，该函数从您的包中调用该函数。在调用该函数时，将电子邮件作为参数传递。您可以在 DAG 级别进一步细化它以仅传递电子邮件所需的信息。

from package import outer_task_success_callback
email = 'xyz@example.com'

def task_success_alert(context):
    dag_id = context['dag'].dag_id
    task_id = context['task_instance']. task_id
    outer_task_success_callback(dag_id, task_id, email)
    
default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2019, 6, 13),
    'on_success_callback': task_success_alert,
    'on_failure_callback': task_failure_alert
}

这将允许您在调用包中的函数之前进行自定义。

附带说明一下，气流具有 smtp 电子邮件功能。您可以利用这些解决方案，而不是编写自己的解决方案。

Answer 2

您可以使用partial创建带有预定义参数的函数，例如：

from functools import partial
new_task_success_alert = partial(task_success_alert, email='your_email')

然后添加新函数作为回调：

on_success_callback=new_task_success_alert