如何让Airflow全局自定义宏隐式获取上下文？

我明白你的需求——就是想让自己写的自定义宏和Airflow内置的ds、execution_date一样，不用手动传参就能直接在模板里用对吧？其实核心是让你的宏能自动拿到Airflow的模板上下文，不用显式传递execution_date之类的参数。下面给你具体的实现步骤：

一、修改宏函数：用get_current_context()隐式获取上下文 #

Airflow 2.x提供了get_current_context()工具函数，它能自动从当前模板渲染环境中抓取上下文，这样你的宏函数就不需要再接收**context参数了。修改后的代码如下：

from airflow.operators.python import get_current_context

def ds_ny():
    """返回纽约时区ds的自定义宏"""
    # 自动获取当前模板上下文
    context = get_current_context()
    execution_date = context.get("execution_date")
    return execution_date.in_timezone("America/New_York").format("%Y-%m-%d")

二、注册你的自定义宏 #

接下来要把这个宏注册到Airflow的宏体系里，有两种方式可选：

方式1：全局注册（所有DAG都能用） #

通过Airflow插件实现全局注册，这样所有DAG都能直接调用macros.ds_ny。在你的Airflow插件目录（默认是$AIRFLOW_HOME/plugins）下创建一个custom_macros_plugin.py文件：

from airflow.plugins_manager import AirflowPlugin
from airflow.operators.python import get_current_context

def ds_ny():
    context = get_current_context()
    return context.get("execution_date").in_timezone("America/New_York").format("%Y-%m-%d")

class CustomMacrosPlugin(AirflowPlugin):
    name = "custom_macros"  # 插件名称，随便起
    macros = [ds_ny]  # 把你的宏加入插件的宏列表

重启Airflow服务后，这个宏就全局可用了。

方式2：DAG级别注册（仅当前DAG可用） #

如果只想在某个特定DAG里用这个宏，直接在DAG定义时通过user_defined_macros参数注册：

from airflow import DAG
from airflow.operators.bash import BashOperator
from datetime import datetime
from airflow.operators.python import get_current_context

def ds_ny():
    context = get_current_context()
    return context.get("execution_date").in_timezone("America/New_York").format("%Y-%m-%d")

with DAG(
    dag_id="test_ny_ds_macro",
    start_date=datetime(2024, 1, 1),
    schedule_interval="@daily",
    user_defined_macros={"ds_ny": ds_ny}  # 注册当前DAG可用的宏
) as dag:
    test_task = BashOperator(
        task_id="print_ny_date",
        bash_command='echo "纽约时区日期：{{ macros.ds_ny }}"'
    )

三、验证效果 #

现在你就可以像使用内置宏一样，直接在模板里写{{ macros.ds_ny }}了，完全不用传任何参数。比如上面的BashOperator任务执行后，会输出对应execution_date转换为纽约时区的日期字符串。

注意事项 #

• 这个方法依赖Airflow 2.x版本，get_current_context()是2.x才引入的API；如果是1.x版本，你需要通过Jinja2的上下文对象来获取，不过更建议升级到2.x版本。
• 全局注册插件后，记得重启Airflow的webserver和scheduler服务，否则插件不会生效。

内容的提问来源于stack exchange，提问作者Yoni M