如何查询数据库获取数据以构建并调度并行Airflow运算符

Airflow的DAG解析阶段（调度器加载DAG代码的阶段）和任务运行阶段是完全独立的时间线：DAG结构必须在解析阶段确定，而XCom、数据库查询结果等都是任务运行阶段才会生成的数据，因此不可能直接把运行时的输出传回DAG解析阶段修改DAG结构，这是Airflow的架构设计决定的。

以下是三种符合Airflow最佳实践的实现方案：

方案1：动态任务映射（优先推荐，Airflow 2.3+版本支持） #

动态任务映射是Airflow官方推出的运行时动态生成任务的特性，完全不需要在DAG顶层执行数据库查询，不会触发DAG加载超时。
实现步骤：

• 定义单独的PythonOperator负责查询MongoDB，返回所有子任务的参数列表，返回结果会自动存入XCom
• 将三阶段的计算逻辑封装为可映射的任务，通过.expand()方法传入上游查询任务的输出，Airflow会在运行时自动根据参数列表的长度生成对应数量的任务
  示例代码：

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime

# 查询MongoDB获取所有子任务参数
def query_mongo_params(**context):
    # 此处写入你的MongoDB查询逻辑，返回结构示例如下
    return [
        {"entity_id": 1, "start_date": "2024-01-01", "end_date": "2024-01-10"},
        {"entity_id": 1, "start_date": "2024-01-11", "end_date": "2024-01-20"},
        # 其余子任务参数
    ]

# 三阶段计算逻辑
def stage1(entity_id, start_date, end_date, **context):
    # 第一阶段业务逻辑
    pass
def stage2(entity_id, start_date, end_date, **context):
    # 第二阶段业务逻辑
    pass
def stage3(entity_id, start_date, end_date, **context):
    # 第三阶段业务逻辑
    pass

with DAG(
    dag_id="multi_entity_process",
    start_date=datetime(2024,1,1),
    schedule_interval="@daily",
    catchup=False
) as dag:
    get_params = PythonOperator(
        task_id="get_mongo_params",
        python_callable=query_mongo_params
    )

    # 动态映射生成所有子任务
    s1 = PythonOperator.partial(task_id="stage1", python_callable=stage1).expand(op_kwargs=get_params.output)
    s2 = PythonOperator.partial(task_id="stage2", python_callable=stage2).expand(op_kwargs=get_params.output)
    s3 = PythonOperator.partial(task_id="stage3", python_callable=stage3).expand(op_kwargs=get_params.output)

    get_params >> s1 >> s2 >> s3

方案2：旧版本Airflow兼容方案（2.3版本以下） #

如果使用的是不支持动态任务映射的旧版本Airflow，可以用固定占位任务+运行时过滤的方式实现：

• 预先估算最大需要的子任务数量，例如200个实体最多每个拆10个日期段，就预留2000个任务槽，在DAG顶层循环生成固定数量的任务组，不会触发数据库查询，不会超时
• 每个任务组内先执行ShortCircuitOperator，读取上游查询MongoDB任务存入XCom的参数列表，判断当前任务组的序号是否在参数列表长度范围内，超出则直接跳过后续任务
• 序号符合要求的任务组，取参数列表中对应位置的参数执行三阶段计算逻辑

方案3：预生成参数配置方案 #

适用于实体的日期范围更新频率较低、不需要每次运行DAG都实时查询的场景：

• 单独维护一个轻量的定时任务（可以是独立crontab，也可以是另一个无 heavy 逻辑的Airflow DAG），定期查询MongoDB生成所有子任务的参数列表，存入JSON格式的配置文件，放在Airflow调度器和工作节点都能读取到的路径
• 业务DAG的顶层代码只读取本地的静态配置文件，无需查询数据库，不会触发加载超时，直接根据配置文件生成对应数量的三阶段任务链即可

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