Airflow动态创建DatabricksSubmitRunOperator任务报错求助

问题背景 #

需要创建Airflow DAG实现历史数据加载自动化，传入jar_path、main_class_name、start_param、end_param、max_run_range5个参数，基于后三者生成日期列表，为每个日期触发DatabricksSubmitRunOperator任务，但执行时出现动态任务错误。

原代码及错误信息 #

原DAG代码 #

@dag(default_args=runParams,
     schedule_interval=None,
     start_date=pendulum.datetime(2024, 4, 22, tz="US/Pacific"),
     tags=["historical_load"]
     )
def CAN_HIST_LOAD_TEST():

    @task
    def get_args():
        context = get_current_context()
        jar_path = context['dag_run'].conf['jar_path']
        main_class_name = context['dag_run'].conf['main_class_name']
        start_param = context['dag_run'].conf['start_param']
        end_param = context['dag_run'].conf['end_param']
        max_run_range = context['dag_run'].conf['max_run_range']

        return {
            'jar_path': jar_path,
            'main_class_name': main_class_name,
            'start_param': start_param,
            'end_param': end_param,
            'max_run_range': max_run_range
        }

    @task
    def process_args(args) -> list[dict]:
        start_date = args['start_param']
        end_date = args['end_param']
        interval = args['max_run_range']

        date_params = []
        idx = 0
        while start_date <= end_date:
            dict_item = {idx: start_date}
            date_params.append(dict_item)
            idx += 1
            start_date = datetime.strptime(start_date, '%Y-%m-%d').date() + timedelta(days=interval)
            start_date = str(start_date)

        args['date_params_dict'] = date_params

        return date_params

    @task
    def spark_task(args, dt_dict):
        dt = ""
        for key, val in dt_dict.items():
            dt = val
        DatabricksSubmitRunOperator(
            task_id=f"process_task_{dt}",
            new_cluster=jobCluster,
            dag=dag,
            databricks_conn_id="${DBxConnID}",
            access_control_list=accessControlList,
            spark_jar_task={"main_class_name": args['main_class_name'], "parameters": dt},
            libraries=[{"jar": args['jar_path']}],
        )

    spark_task.partial(args=get_args()).expand(dt_dict=process_args(get_args()))

CAN_HIST_LOAD_TEST_DAG = CAN_HIST_LOAD_TEST()

错误信息 #

Traceback (most recent call last):
  File "/usr/local/lib/python3.9/site-packages/airflow/decorators/base.py", line 188, in execute
    return_value = super().execute(context)
  File "/usr/local/lib/python3.9/site-packages/airflow/operators/python.py", line 175, in execute
    return_value = self.execute_callable()
  File "/usr/local/lib/python3.9/site-packages/airflow/operators/python.py", line 193, in execute_callable
    return self.python_callable(*self.op_args, **self.op_kwargs)
  File "/usr/local/airflow/dags/can/CAN_HIST_LOAD_TEST.py", line 122, in spark_task
    DatabricksSubmitRunOperator(
  File "/usr/local/lib/python3.9/site-packages/airflow/models/baseoperator.py", line 376, in apply_defaults
    task_group = TaskGroupContext.get_current_task_group(dag)
  File "/usr/local/lib/python3.9/site-packages/airflow/utils/task_group.py", line 489, in get_current_task_group
    return dag.task_group
AttributeError: 'function' object has no attribute 'task_group'

问题分析 #

1. 错误核心原因：在@task装饰的Python任务函数内部实例化DatabricksSubmitRunOperator，并且传入了dag=dag——这里的dag是DAG定义函数本身（CAN_HIST_LOAD_TEST），而非实际的DAG实例，导致Airflow无法识别任务所属的DAG上下文。
2. 参数格式冗余：生成日期参数时用字典包裹，增加了解析成本，直接返回日期字符串列表更适合动态任务映射。
3. Operator参数错误：spark_jar_task的parameters要求是列表类型，原代码传入单个字符串会导致参数解析失败。

修复方案 #

1. 直接对DatabricksSubmitRunOperator使用partial+expand实现动态任务映射，无需在Python任务内部创建Operator。
2. 简化日期参数生成逻辑，返回字符串列表而非字典。
3. 移除dag=dag参数，DAG上下文内创建的Operator会自动关联当前DAG。
4. 修正spark_jar_task的parameters格式为列表。

修复后的完整代码 #

from airflow.decorators import dag, task
from airflow.providers.databricks.operators.databricks import DatabricksSubmitRunOperator
from airflow.utils.context import get_current_context
import pendulum
from datetime import datetime, timedelta

# 假设runParams、jobCluster、accessControlList已提前定义
runParams = {}
jobCluster = {}
accessControlList = []

@dag(default_args=runParams,
     schedule_interval=None,
     start_date=pendulum.datetime(2024, 4, 22, tz="US/Pacific"),
     tags=["historical_load"]
     )
def CAN_HIST_LOAD_TEST():

    @task
    def get_args():
        context = get_current_context()
        conf = context['dag_run'].conf
        return {
            'jar_path': conf['jar_path'],
            'main_class_name': conf['main_class_name'],
            'start_param': conf['start_param'],
            'end_param': conf['end_param'],
            'max_run_range': int(conf['max_run_range'])  # 转为整数避免后续计算错误
        }

    @task
    def generate_date_range(args) -> list[str]:
        start_date = datetime.strptime(args['start_param'], '%Y-%m-%d').date()
        end_date = datetime.strptime(args['end_param'], '%Y-%m-%d').date()
        interval_days = args['max_run_range']

        date_list = []
        current_date = start_date
        while current_date <= end_date:
            date_list.append(current_date.strftime('%Y-%m-%d'))
            current_date += timedelta(days=interval_days)
        
        return date_list

    # 获取参数和日期列表
    args = get_args()
    date_range = generate_date_range(args)

    # 动态生成Databricks任务
    DatabricksSubmitRunOperator.partial(
        new_cluster=jobCluster,
        databricks_conn_id="${DBxConnID}",
        access_control_list=accessControlList,
        libraries=[{"jar": args['jar_path']}],
    ).expand(
        task_id=[f"process_task_{dt}" for dt in date_range],
        spark_jar_task=[{"main_class_name": args['main_class_name'], "parameters": [dt]} for dt in date_range]
    )

CAN_HIST_LOAD_TEST_DAG = CAN_HIST_LOAD_TEST()

关键说明 #

• 用DatabricksSubmitRunOperator.partial()固定公共参数，expand()动态传入每个任务的唯一参数（task_id和spark_jar_task）。
• 日期生成逻辑直接返回字符串列表，简化动态映射的参数处理。
• 确保max_run_range转为整数，避免日期计算时的类型错误。
• spark_jar_task的parameters改为列表格式，符合Databricks Operator的参数要求。

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