如何在自定义Operator中解析dag_run.conf至自定义函数？遇模板错误求解

代码示例 #

def render_file_name(task_type, table_name):

    if 'cleanup' == task_type:
        file_name = f"path/to/{table_name}_his_full_cleanup.sql"
    else:
        raise ValueError(f"Unknown task type: {task_type}")

    return file_name

with DAG("load_data", template_searchpath=SCRIPTS_PATH) as dag:

    cleanup_sql = render_file_name(task_type='cleanup', table_name='{{ dag_run.conf["table_name"] }}')

    table_his_full_cleanup = MyCustomOperator(
    task_id="table_full_cleanup",
    name="table_full_cleanup",
    sql=cleanup_sql,
    parameters={
        "env": ENV,
        "temp_table_name": '{{ dag_run.conf["table_name"] }}_temp',
    },
    default_iam_role=IAM_ROLE,
    spark_cluster=MY_SPARK,
    )

运行错误 #

jinja2.exceptions.TemplateNotFound: path/to/{{ dag_run.conf["table_name"] }}_his_full_cleanup.sql

疑问 #

sql参数无法解析dag_run.conf["table_name"]，但parameters参数能正确解析该值，请问原因是什么？该如何解决？

1. Python函数执行时机问题：render_file_name是普通Python函数，在Airflow解析DAG定义的阶段就会执行，此时{{ dag_run.conf["table_name"] }}只是一个字符串字面量，不会被Jinja2解析，直接拼接到文件名中，导致Airflow尝试寻找包含模板语法的文件，自然找不到。
2. 模板字段配置差异：Airflow的Operator会对标记为template_fields的参数进行Jinja2模板渲染，parameters是多数Operator默认的模板字段，会在任务运行阶段解析变量；而MyCustomOperator大概率没有将sql字段加入template_fields，所以sql参数中的模板语法不会被解析，直接以原始字符串传递。

方法1：修改MyCustomOperator，将sql加入模板字段 #

如果可以修改Operator源码，在MyCustomOperator类中添加template_fields属性，把sql包含进去：

class MyCustomOperator(BaseOperator):
    template_fields = ("sql", "parameters")  # 确保sql在模板字段列表内
    # 其他原有代码逻辑

修改后，Airflow会自动在任务运行阶段解析sql中的模板语法，此时可以直接在sql参数中写模板表达式，无需提前用Python函数拼接：

with DAG("load_data", template_searchpath=SCRIPTS_PATH) as dag:
    table_his_full_cleanup = MyCustomOperator(
        task_id="table_full_cleanup",
        name="table_full_cleanup",
        sql="path/to/{{ dag_run.conf['table_name'] }}_his_full_cleanup.sql",
        parameters={
            "env": ENV,
            "temp_table_name": '{{ dag_run.conf["table_name"] }}_temp',
        },
        default_iam_role=IAM_ROLE,
        spark_cluster=MY_SPARK,
    )

方法2：直接用Jinja2语法拼接文件名（无需修改Operator） #

如果无法修改Operator，可直接在sql参数中用Jinja2表达式完成文件名拼接，前提是sql字段本身支持模板渲染（若Operator未配置template_fields则此方法无效）：

with DAG("load_data", template_searchpath=SCRIPTS_PATH) as dag:
    table_his_full_cleanup = MyCustomOperator(
        task_id="table_full_cleanup",
        name="table_full_cleanup",
        sql="{{ 'path/to/' ~ dag_run.conf['table_name'] ~ '_his_full_cleanup.sql' }}",
        parameters={
            "env": ENV,
            "temp_table_name": '{{ dag_run.conf["table_name"] }}_temp',
        },
        default_iam_role=IAM_ROLE,
        spark_cluster=MY_SPARK,
    )

注：Jinja2中字符串拼接用~运算符，避免Python字符串格式化的时机问题。

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