Apache Airflow DAG代码自动生成，使用Springboot+Freemarker

Apache Airflow DAG代码自动生成，使用Springboot+Freemarker

---

前言

目前，Apache Airflow并没有直接支持从图形流程图生成可用的DAG Python代码的官方工具或库。因此，如果要根据自己画的流程图生成Airflow的DAG代码，只能考虑其他手动或自定义的方法，如自定义脚本、模板引擎、或者一些转换工具。

自动生成Airflow DAG代码可能涉及到很多复杂的问题，因为需要将图形元素（节点、边等）映射到Airflow的概念（任务、依赖关系等）。在实际实现中，可能需要一些自定义逻辑和调整以确保生成的代码是正确且符合Airflow的要求的。实际应用中我们打算使用Freemarker模板引擎来实现DAG代码自动生成。

---

一、用途？

我们实现的拖拽组件需要有流程控制，如图，

Airflow自身没有自动生成流程代码的能力，需要额外开发自动生成流程代码的功能。Freemarker Java模板引擎具有强大的表达式语言和广泛的用例。它可以与Spring Boot集成，并支持在模板中使用FreeMarker的语法。
在画好流程图后，保存即可生成Airflow的python代码，并加载到airflow中，以便执行整个流程。

二、设计思路

使用Freemarker结合自定义标记语法，将流程图中的信息注入到DAG代码模板中：

1. 创建流程图:
   首先，需要创建一个流程图，我们自己开发了流程图绘制画布画板。

2. 定义自定义标记:
   在流程图上为节点和边定义自定义标记，我们的节点已经做好预置的python功能脚本，并设置了特定格式的文本来表示任务的ID、参数等信息。包括类似{{ task_id }}这样的标记。

3. 通过FreeMarker模板引擎生成Apache Airflow的Python代码，并将生成的代码保存到Airflow的DAG目录中。

4. 示例代码如下：

DagCodeGenerator.java

用于生成DAG代码文件的工具类，该类负责将生成的Python代码写入文件。

import org.springframework.stereotype.Component;

import java.io.FileWriter;
import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;

@Component
public class DagCodeGenerator {

    public void generateAndSaveDagCode(String dagCode, String dagId) throws IOException {
        // 定义Airflow DAG目录路径
        String airflowDagDirectory = "/home/airflow/dags";

        // 构建Python文件路径
        Path pythonFilePath = Paths.get(airflowDagDirectory, dagId + ".py");

        // 将生成的Python代码写入文件
        try (FileWriter writer = new FileWriter(pythonFilePath.toFile())) {
            writer.write(dagCode);
        }
    }
}

DagController.java

该Controller负责渲染FreeMarker模板并返回生成的Airflow DAG代码

import freemarker.template.Configuration;
import freemarker.template.Template;
import freemarker.template.TemplateException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Controller;
import org.springframework.ui.Model;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.ResponseBody;

import java.io.IOException;
import java.io.StringWriter;
import java.util.HashMap;
import java.util.Map;

@Controller
@RequestMapping("/dag")
public class DagController {

    @Autowired
    private Configuration freemarkerConfig;

    @Autowired
    private DagCodeGenerator dagCodeGenerator;

    @GetMapping("/generate")
    @ResponseBody
    public String generateDag(Model model) throws IOException, TemplateException {
        // 从文件加载FreeMarker模板
        Template template = freemarkerConfig.getTemplate("dag_template.ftl");

        // 定义变量
        Map<String, Object> dataModel = new HashMap<>();
        dataModel.put("dag_id", "my_dag");
        dataModel.put("schedule_interval", "daily");
        dataModel.put("tasks", List.of(Map.of("task_id", "task1"), Map.of("task_id", "task2")));
        dataModel.put("dependencies", List.of(Map.of("upstream", "task1", "downstream", "task2")));

        // 渲染模板
        StringWriter writer = new StringWriter();
        template.process(dataModel, writer);

        // 获取生成的DAG代码
        String dagCode = writer.toString();

        // 生成并保存DAG代码到文件
        dagCodeGenerator.generateAndSaveDagCode(dagCode, "my_dag");

        // 返回成功消息或其他响应
        return "DAG code generated and saved successfully!";
    }
}

application.properties

spring.freemarker.template-loader-path=classpath:/templates/

dag_template.ftl

FreeMarker模板文件，用于生成Airflow DAG的Python代码。

# dag_template.ftl
from airflow import DAG
from airflow.operators.dummy_operator import DummyOperator

dag = DAG('{{ dag_id }}', schedule_interval='{{ schedule_interval }}')

{% for task_info in tasks %}
{{ task_info['task_id'] }} = DummyOperator(task_id='{{ task_info['task_id'] }}', dag=dag)
{% endfor %}

{% for dependency in dependencies %}
{{ dependency['upstream'] }} >> {{ dependency['downstream'] }}
{% endfor %}

通过访问 /dag/generate 路径，将触发生成Airflow DAG代码并将其保存到文件。

---

总结

以上是个引子，为后续我们生成各种DAG代码提供一个思路，前置条件是每个节点和连线的python模块代码已经实现。