如何从BranchPythonOperator调用含动态子任务ID的TaskGroup？

业务需求 #

检查BigQuery中指定表是否存在：

• 若存在，直接结束DAG；
• 若不存在，执行从Postgres到Google Cloud Storage的数据导入操作。

实现思路 #

通过BranchPythonOperator做分支判断，调用包含动态生成子任务的TaskGroup完成数据导入。子任务ID由循环生成（如parent_task_id.sub_task_1、parent_task_id.sub_task_2等）。

错误信息 #

运行DAG时触发：

airflow.exceptions.TaskNotFound: Task parent_task_id not found

当前代码 #

DAG代码 #

import airflow
from airflow.providers.google.cloud.transfers.postgres_to_gcs import PostgresToGCSOperator
from airflow.utils.task_group import TaskGroup
from google.cloud.exceptions import NotFound
from airflow import DAG
from airflow.operators.python import BranchPythonOperator
from airflow.operators.dummy import DummyOperator
from google.cloud import bigquery

default_args = {
    'owner': 'Airflow',
    'start_date': airflow.utils.dates.days_ago(2),
}

with DAG(dag_id='branch_dag', default_args=default_args, schedule_interval=None) as dag:

    def create_task_group(worker=1):
        var  = dict()
        with TaskGroup(group_id='parent_task_id') as tg1:
            for i in range(worker):
                var[f'sub_task_{i}'] = PostgresToGCSOperator(
                    task_id = f'sub_task_{i}',
                    postgres_conn_id = 'some_postgres_conn_id',
                    sql = 'test.sql',
                    bucket = 'test_bucket',
                    filename = 'test_file.json',
                    export_format = 'json',
                    gzip = True,
                    params = {
                        'worker': worker
                    }
                )
        return tg1
    
    def is_exists_table():
        client = bigquery.Client()
        try:
            table_name = client.get_table('dataset_id.some_table')
            if table_name:
                return 'task_end'
        except NotFound as error:       
            return 'parent_task_id'

    task_start = DummyOperator(
        task_id = 'start'
        )

    task_branch_table = BranchPythonOperator(
        task_id ='check_table_exists_in_bigquery',
        python_callable = is_exists_table
        )

    task_pg_to_gcs_init = create_task_group(worker=3)

    task_end = DummyOperator(
        task_id = 'end',
        trigger_rule = 'all_done'
    )    

    task_start >> task_branch_table >> task_end
    task_start >> task_branch_table >> task_pg_to_gcs_init >> task_end

SQL脚本（test.sql） #

SELECT 
id,
name,
country
FROM some_table
WHERE 1=1
AND ABS(MOD(hashtext(id::TEXT), 3)) = {{params.worker}};

-- returns 1M+ rows

错误核心原因 #

BranchPythonOperator要求返回具体的可执行任务ID，而你返回的parent_task_id是TaskGroup的容器ID，TaskGroup本身不是可执行任务，Airflow无法找到对应的任务实例，因此触发报错。

以下提供两种可行的解决方式：

方案1：返回TaskGroup内所有子任务ID列表 #

BranchPythonOperator支持返回任务ID列表，直接返回所有动态生成的子任务完整ID（格式为group_id.sub_task_id）即可触发整个TaskGroup的执行。

修改代码如下：

def create_task_group(worker=1):
    sub_task_ids = []
    with TaskGroup(group_id='parent_task_id') as tg1:
        for i in range(worker):
            task_id = f'sub_task_{i}'
            # 记录完整子任务ID
            sub_task_ids.append(f'parent_task_id.{task_id}')
            PostgresToGCSOperator(
                task_id = task_id,
                postgres_conn_id = 'some_postgres_conn_id',
                sql = 'test.sql',
                bucket = 'test_bucket',
                filename = f'test_file_{i}.json',  # 每个子任务文件名需唯一，避免数据覆盖
                export_format = 'json',
                gzip = True,
                params = {
                    'worker': i  # 原代码传worker会导致所有子任务过滤条件相同，改为传i
                }
            )
    return tg1, sub_task_ids

# 获取TaskGroup实例和子任务ID列表
task_pg_to_gcs_init, pg_sub_task_ids = create_task_group(worker=3)

def is_exists_table():
    client = bigquery.Client()
    try:
        client.get_table('dataset_id.some_table')
        return 'end'
    except NotFound:       
        return pg_sub_task_ids  # 返回所有子任务ID列表

方案2：给TaskGroup添加入口Dummy任务（更简洁） #

在TaskGroup内新增一个入口Dummy任务，分支函数返回该入口任务的ID，由入口任务触发所有子任务。这种方式无需维护子任务ID列表，更易维护。

修改代码如下：

def create_task_group(worker=1):
    with TaskGroup(group_id='parent_task_id') as tg1:
        # 添加TaskGroup入口任务
        tg_start = DummyOperator(task_id='tg_start')
        sub_tasks = []
        for i in range(worker):
            task = PostgresToGCSOperator(
                task_id = f'sub_task_{i}',
                postgres_conn_id = 'some_postgres_conn_id',
                sql = 'test.sql',
                bucket = 'test_bucket',
                filename = f'test_file_{i}.json',
                export_format = 'json',
                gzip = True,
                params = {
                    'worker': i
                }
            )
            sub_tasks.append(task)
        # 设置依赖：入口任务触发所有子任务
        tg_start >> sub_tasks
    return tg1

def is_exists_table():
    client = bigquery.Client()
    try:
        client.get_table('dataset_id.some_table')
        return 'end'
    except NotFound:       
        return 'parent_task_id.tg_start'  # 返回TaskGroup入口任务的完整ID

最终依赖关系 #

两种方案的DAG依赖关系保持不变：

task_start >> task_branch_table >> task_end
task_start >> task_branch_table >> task_pg_to_gcs_init >> task_end

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