python - 在 Airflow 中，如何使用上下文将参数传递给 on_success_callback 函数处理程序？

Question

在 Airflow 中，如何使用上下文将参数传递给 on_success_callback 函数处理程序？

这是我的测试代码:

import airflow
from airflow import DAG
from airflow.operators import MSTeamsWebhookOperator
from airflow.operators.bash_operator import BashOperator
from airflow.operators.dummy_operator import DummyOperator
from datetime import datetime
from transaction_analytics import helpers
from airflow.utils.helpers import chain

# Parameters & variables
schedule_interval        = "0 20 * * *"

def _task_success_callback(context):
    dagid         = context["task_instance"].dag_id
    duration      = context["task_instance"].duration
    executiondate = context["execution_date"]
    logurl        = context["task_instance"].log_url.replace("localhost", "agbqhsbldd017v.agb.rbxd.ds")# workaround until we config airflow
    pp1           = context["params"].param1 
    #pp1          = "{{ params.param1 }}"
    ms_teams_op = MSTeamsWebhookOperator(
        task_id="success_notification",
        http_conn_id="msteams_airflow",
        message="DAG {ppram1} `{dag}`  finished successfully!".format(dag=context["task_instance"].dag_id, ppram1=pp1),
        subtitle="Execution Date = {p1}, Duration = {p2}".format(p1=executiondate,p2=duration),
        button_text = "View log",
        button_url = "{log}".format(log=logurl),
        theme_color="00FF00"#,
        #proxy= "http://10.72.128.202:3128"
    )
    ms_teams_op.execute(context)

    main_dag = DAG('test_foley',
            schedule_interval=schedule_interval,
            description='Test foley',
            start_date=datetime(2020, 4, 19),
            default_args=None,
            max_active_runs=2,
            default_view='graph',     # Default view graph
            #orientation='TB', # Top-Bottom graph
            on_success_callback=_task_success_callback,
            #on_failure_callback=outer_task_failure_callback,
            catchup=False, # Do not catchup, run only latest
            params={
              "param1": "value1",
              "param2": "value2"
            }
           )

################################### START ######################################
dag_chain = []

start = DummyOperator(task_id='start', retries = 3, dag=main_dag)
dag_chain.append(start)

step1 = BashOperator(
    task_id='step1',
    bash_command='pwd',
    dag=main_dag,
)
dag_chain.append(step1)

step2 = BashOperator(
    task_id='step2',
    bash_command='exit 0',
    dag=main_dag,
)
dag_chain.append(step2)


end = DummyOperator(task_id='end', dag=main_dag)
dag_chain.append(end)

chain(*dag_chain)

我有一个处理成功的事件处理函数 _task_success_callback。在 DAG 中，我有捕获该事件的 on_success_callback=_task_success_callback。

它有效...但现在我需要将一些参数传递给 _task_success_callback。什么是最好的方法？

当该函数接收上下文时，我尝试在 DAG 中创建参数，如您所见:

        params={
          "param1": "value1",
          "param2": "value2"
        }

但我似乎无法访问它们？

我的问题是:

1. 访问参数我做错了什么？
2. 有没有更好的参数传递方式？

注意:我看到了这个类似的问题 How to pass parameters to Airflow on_success_callback and on_failure_callback有一个答案......并且有效。但是我正在寻找的是使用上下文传递参数....

Answer 1

回想一下，Airflow 进程文件只是 Python，如果您在解析过程中没有引入太多开销(因为 Airflow 频繁解析文件，并且开销会加起来)，您可以使用 Python 可以做的一切。特别是对于您的情况，我建议为您的回调返回一个嵌套函数(闭包):

将它放在与您的 Airflow 进程相邻的文件中，比方说 on_callbacks.py

def success_ms_teams(param_1, param_2):

    def callback_func(context):
        print(f"param_1: {param_1}")
        print(f"param_2: {param_2}")
        # ... trimmed for brevity ...#
        ms_teams_op.execute(context)

    return callback_func

然后在您的流程中您可以这样做:

from airflow import models

from on_callbacks import success_ms_teams

with models.DAG(
    ...
    on_success_callback=success_ms_teams(
        "value1", # These values become the
        "value2", # `param_1` and `param_2` 
    )
) as dag:
    ...