Airflow Task 血缘分析

血缘分析流程

Airflow 原生 CLI 与 REST API 服务，Airflow Task 血缘方式可以通过以下几个步骤获取：

1. 通过 CLI: airflow list_dags 方式获取 dag 列表 list_dags
2. 遍历 dag 列表，通过 CLI: airflow show_dag ${dag_id} 方式获取 dag 描述信息，标红出即为 dag 内 task 依赖关系 show_dags
3. 由于依赖关系存在 ExternalTaskSensor，所以需要通过 REST API 获取起 Task Info 信息 rest

4. 通过绑定关系解析，生成 Parent - Child 关系即可

核心代码实现

List Dags

        String command = "airflow list_dags";
        String result = CommandUtils.exeCommand(command);
        // 解析返回数据类型
        List<String> shellResultLines = CommandUtils.getShellResultLines(result);
        int lastIndex = 0;
        for(int i=0; i<shellResultLines.size(); i++) {
            if (shellResultLines.get(i).contains("--------")) {
                lastIndex = i;
            }
        }
        List<String> dagList = shellResultLines.subList(lastIndex + 1, shellResultLines.size());
        return dagList;

Show Dag

        String command = "airflow show_dag %s";

        // check 有效 dag 列表
        List<String> validDagList = dagList.stream().map(dag -> {
            String showDagCommand = String.format(command, dag);
            try {
                int exeCode = CommandUtils.exeCommand(showDagCommand, null);
                if (exeCode == 0) {
                    return dag;
                }
            } catch (IOException e) {
                log.error("exe command: {} is error, Exception:{} ", showDagCommand);
            }
            return null;
        }).filter(dag -> dag != null).collect(Collectors.toList());

        // 获取有效 dag 信息
        List<Map<String, String>> dagDetailList = validDagList.stream().map(dag -> {
            String result = CommandUtils.exeCommand(String.format(command, dag));
            HashMap<String, String> dagDetailMap = new HashMap();
            dagDetailMap.put(dag, result);
            return dagDetailMap;
        }).collect(Collectors.toList());

Task

由于部分单独包含单个 Task 的 Dag，在 Show Dag 无法获取，所以添加无依赖关系 Task

// 提取 task 关系
        dagDetailList.stream().forEach(dagDetailMap -> {
            List<DagTaskDependencyDO> dagTaskDependencyList = new ArrayList<>();
            dagDetailMap.forEach((dag, detail) -> {
                // 拿到 dag 下 所有 tasks
                List<String> taskList = listTasks(dag);

                List<String> shellResultLines = CommandUtils.getShellResultLines(detail);
                Set<String> relationTaskSet = new HashSet<>();
                shellResultLines.forEach(line -> {

                    // 检查包含的依赖关系
                    if (line.contains("->")) {
                        String[] taskIds = line.split("->");
                        String parentTask = taskIds[0].trim().replaceAll("\"|\'", "");
                        String task = taskIds[1].trim().replaceAll("\"|\'", "");
                        DagTaskDependencyDO dagTaskDependencyDO = new DagTaskDependencyDO();
                        dagTaskDependencyDO.setDagId(dag);
                        dagTaskDependencyDO.setTaskId(task);
                        dagTaskDependencyDO.setParentDagId(dag);
                        dagTaskDependencyDO.setParentTaskId(parentTask);
                        dagTaskDependencyList.add(dagTaskDependencyDO);
                        relationTaskSet.add(task);
                        relationTaskSet.add(parentTask);
                    }
                });
                // 添加无依赖关系的 task
                taskList.removeAll(relationTaskSet);
                taskList.forEach(task -> {
                    DagTaskDependencyDO dagTaskDependencyDO = new DagTaskDependencyDO();
                    dagTaskDependencyDO.setDagId(dag);
                    dagTaskDependencyDO.setTaskId(task);
                    dagTaskDependencyDO.setParentDagId("-");
                    dagTaskDependencyDO.setParentTaskId("-");
                    dagTaskDependencyList.add(dagTaskDependencyDO);
                });
            });

External Sensor

1. 部分 External Sensor 对应的 Dag Task，需要替换映射回 External Dag && Task
2. ExternalTaskMarker 需要特殊处理，详见下述代码

                List<DagTaskDependencyDO> suitableDagDetailList = dagTaskDependencyList.stream().map(dagTaskDependency -> {
                DagTaskDependencyDO target = new DagTaskDependencyDO();
                try {
                    BeanUtils.copyProperties(target, dagTaskDependency);
                } catch (IllegalAccessException e) {
                    e.printStackTrace();
                } catch (InvocationTargetException e) {
                    e.printStackTrace();
                }

                String response = HttpUtils.getRequest(RestApiHelper.getTaskInfoUrl(env, dagTaskDependency.getDagId(), dagTaskDependency.getTaskId()));
                Map<String, String> externalMap = wideExternal(response);
                if (!externalMap.isEmpty()) {
                    if (! target.getParentDagId().equals("-")) {
                        target.setRawDagId(dagTaskDependency.getDagId());
                        target.setRawTaskId(dagTaskDependency.getTaskId());
                        short isExternalSensor = 1;
                        target.setIsExternalSensor(isExternalSensor);
                        target.setDagId(externalMap.get(RestResponseConstant.EXTERNAL_DAG_ID));
                        target.setTaskId(externalMap.get(RestResponseConstant.EXTERNAL_TASK_ID));
                    }
                    // 如果存在单个 task，且配置 ExternalTaskMarker，则保存 parentDagId && parentTaskId
                    if (target.getParentDagId().equals("-")) {
                        short isExternalSensor = 1;
                        target.setIsExternalSensor(isExternalSensor);
                        target.setParentDagId(externalMap.get(RestResponseConstant.EXTERNAL_DAG_ID));
                        target.setParentTaskId(externalMap.get(RestResponseConstant.EXTERNAL_TASK_ID));
                    }

                }
                String parentResponse = HttpUtils.getRequest(RestApiHelper.getTaskInfoUrl(env, dagTaskDependency.getParentDagId(), dagTaskDependency.getParentTaskId()));
                Map<String, String> parentExternalMap = wideExternal(parentResponse);
                if (!parentExternalMap.isEmpty()) {
                    target.setRawParentDagId(dagTaskDependency.getParentDagId());
                    target.setRawParentTaskId(dagTaskDependency.getParentTaskId());
                    short isExternalSensor = 1;
                    target.setIsExternalSensor(isExternalSensor);
                    target.setParentDagId(parentExternalMap.get(RestResponseConstant.EXTERNAL_DAG_ID));
                    target.setParentTaskId(parentExternalMap.get(RestResponseConstant.EXTERNAL_TASK_ID));
                }

                return target;
            }).collect(Collectors.toList());
            batchInsert(env, suitableDagDetailList);
        });

结果