使用 dag-factory 简化 Apache Airflow DAG 创建

什么是 dag-factory？

dag-factory 是一个 Python 库，用于动态生成 Apache Airflow DAG。它允许你使用 YAML 文件定义 DAG，而不是编写 Python 代码。这使得 DAG 的定义更加简洁和易读。

dag-factory 的主要功能

1. 减少代码量：使用 YAML 文件定义 DAG，可以减少 Python 代码的编写和维护。
2. 提高可读性：YAML 格式比 Python 代码更易于阅读和理解。
3. 动态生成 DAG：根据配置参数动态生成 DAG，适用于需要创建多个相似 DAG 的场景。
4. 简化 DAG 管理：将 DAG 定义存储在 YAML 文件中，方便版本控制和自动化部署。

使用 dag-factory 的步骤

步骤 1：安装 dag-factory

首先，你需要在你的 Airflow 环境中安装 dag-factory。如果你使用 Astronomer 的 Astro 项目，可以在 requirements.txt 文件中添加以下行：

dag-factory<=1.0.0

然后保存并更新依赖项。

如果不是 Astronomer 用户，可以使用 pip 安装：

pip install dag-factory<=1.0.0

步骤 2：创建 YAML 配置文件

在 Airflow 的 dags 目录下创建一个名为 configs 的子目录，用于存放 YAML 配置文件。例如，创建一个名为 my_dag.yaml 的文件：

my_dag:
  default_args:
    owner: 'example_owner'
    start_date: '2024-07-01'
    end_date: '2030-01-01'
    retries: 1
    retry_delay_sec: 300
    schedule_interval: '0 3 * * *'
    catchup: False
    concurrency: 1
    max_active_runs: 1
    dagrun_timeout_sec: 60

步骤 3：创建 DAG Factory 文件

在 dags 目录下创建一个名为 dag_factory.py 的 Python 文件，用于将 YAML 文件转换为 DAG。内容如下：

from pathlib import Path
from airflow import DAG
from airflow.configuration import conf as airflow_conf
from dagfactory import load_yaml_dags

config_dir = Path(airflow_conf.get("core", "dags_folder")) / "configs"
load_yaml_dags(globals_dict=globals(), dags_folder=config_dir)

示例代码：动态生成 DAG

如果你需要根据不同配置动态生成 DAG，可以使用 Python 脚本来生成 YAML 文件。例如，创建一个名为 generate_dynamic_dag.py 的脚本：

import yaml

# 定义模板 YAML 文件路径
PATH_TO_TEMPLATE_YAML = "include/template.yml"

# 读取模板 YAML 文件
with open(PATH_TO_TEMPLATE_YAML, "r") as yaml_file:
    template = yaml.load(yaml_file, yaml.SafeLoader)

# 替换占位符
template['my_dag']['default_args']['start_date'] = '2024-11-01'
template['my_dag']['tasks']['task_1']['bash_command'] = 'echo "Hello World!"'

# 写入新的 YAML 文件
with open("dags/dynamic_etl.yml", "w") as ingestion_config:
    ingestion_config.write(yaml.dump(template))

模板 YAML 文件 (template.yml) 可以包含占位符，例如：

my_dag:
  default_args:
    owner: 'example_owner'
    start_date: "<< start_date >>"
    retries: 1
    retry_delay_sec: 300
    schedule_interval: "<< schedule >>"
    catchup: False
  tasks:
    task_1:
      operator: airflow.operators.bash_operator.BashOperator
      bash_command: "<< bash_command_task_1 >>"

通过运行 generate_dynamic_dag.py 脚本，你可以根据不同的配置参数生成多个 YAML 文件，从而动态生成 DAG。

使用场景

• 大量相似 DAG：如果需要创建多个结构相似的 DAG，但参数不同，dag-factory 可以帮助你避免重复编写代码。
• 动态生成 DAG：根据外部数据或配置动态生成 DAG。
• 简化 DAG 管理：使用更简洁、更易读的方式定义和管理 DAG。