如何在Airflow的DatabricksSubmitRunOperator中为execution_timeout传Jinja模板？

问题场景 #

在Airflow中，Jinja模板通常放在双引号内传递给算子，但部分字段（比如DatabricksSubmitRunOperator的execution_timeout）要求传入特定类型对象（此处为timedelta），直接使用模板字符串会触发类型错误。

错误写法1：直接传递模板字符串 #

from airflow import DAG
from airflow.providers.databricks.operators.databricks import DatabricksSubmitRunOperator
from datetime import datetime

default_args = {
    'start_date': datetime(2023, 3, 16)
}

with DAG('my_dag', default_args=default_args, schedule_interval=None) as dag:
    
    task1 = DatabricksSubmitRunOperator(
        task_id='task1',
        databricks_conn_id='databricks_default',
        existing_cluster_id="xyz-cluster-id",
        run_name='My Run',
        wait_for_completion=True,
        execution_timeout="{{ ti.xcom_pull(task_ids='previous_task_id', key='my_timeout') }}"
    )

    task2 = ...

    task1 >> task2

报错信息：

ValueError: execution_timeout must be timedelta object but passed as type <class 'str'>

错误写法2：嵌套模板到timedelta构造 #

task1 = DatabricksSubmitRunOperator(
        task_id='task1',
        databricks_conn_id='databricks_default',
        existing_cluster_id="xyz-cluster-id",
        run_name='My Run',
        wait_for_completion=True,
        execution_timeout=timedelta(minutes="{{ ti.xcom_pull(task_ids='previous_task_id', key='my_timeout') }}")
    )

报错信息：

unsupported type for timedelta minutes component: str

问题根源在于execution_timeout并非Airflow默认的模板化字段，无法直接通过Jinja模板动态赋值，必须自定义包装类扩展模板支持，并处理类型转换。

解决方案：自定义算子包装类 #

步骤1：创建支持execution_timeout模板化的自定义算子 #

继承原算子，将目标字段加入模板字段列表，并在执行阶段完成模板渲染与类型转换：

from airflow.providers.databricks.operators.databricks import DatabricksSubmitRunOperator
from datetime import timedelta

class CustomDatabricksSubmitRunOperator(DatabricksSubmitRunOperator):
    # 将execution_timeout追加到模板字段列表
    template_fields = (*DatabricksSubmitRunOperator.template_fields, 'execution_timeout')

    def execute(self, context):
        # 渲染execution_timeout的Jinja模板
        rendered_timeout = self.render_template(self.execution_timeout, context)
        # 将渲染后的分钟数字符串转换为timedelta对象
        if rendered_timeout:
            self.execution_timeout = timedelta(minutes=int(rendered_timeout))
        # 调用父类执行逻辑
        return super().execute(context)

步骤2：在DAG中使用自定义算子 #

from airflow import DAG
from datetime import datetime

default_args = {
    'start_date': datetime(2023, 3, 16)
}

with DAG('my_dag', default_args=default_args, schedule_interval=None) as dag:
    
    task1 = CustomDatabricksSubmitRunOperator(
        task_id='task1',
        databricks_conn_id='databricks_default',
        existing_cluster_id="xyz-cluster-id",
        run_name='My Run',
        wait_for_completion=True,
        # 直接传入Jinja模板字符串
        execution_timeout="{{ ti.xcom_pull(task_ids='previous_task_id', key='my_timeout') }}"
    )

    task2 = ...

    task1 >> task2

关键说明 #

• 核心操作是将目标字段加入算子的template_fields，让Airflow自动处理Jinja模板渲染
• 模板渲染结果默认是字符串，需手动转换为字段要求的类型（如timedelta）
• 该方案适用于所有需要动态赋值非模板化字段的自定义算子场景

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