I have following use case:
I am trying to create an Airflow DAG which will be used to automate historical data load. I am passing 5 arguments to this DAG - jar_path, main_class_name, start_param, end_param, max_run_range.
Based on start_param, end_param and max_run_range, a list of dates is created. For each individual date of this list, I want to trigger a DatabricksSubmitRunOperator task.
I have below code but when I execute it, I am getting this error for the dynamic tasks that it creates -
Airflow DAG code:
@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()
Error:
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'
Even though I am not using TaskGroup, I am not sure why is it giving this error.
Env used: Azure Databricks, Astronomer for Airflow(airflow version 2.4.3), jar is created on spark/scala code.
I have also tried passing the dates list directly into .expand() method but it also resulted in the same error. 
For a sample argument of 3 dates (start_param='2024-03-01',end_param='2024-03-03',max_run_range=1), the code is able to create 3 mapped instances but it fails in each instance with the above error.
(Ignore start and end tasks)
For other variations that I tried, the code failed without even creating dynamic tasks.
