Pythonのastモジュールを活用してAirflowのDAG依存関係をDAG Docsに表示する

{
  "dag_21": {
    "upstream": [
      "dag_11"
    ],
    "downstream": []
  },
  "dag_11": {
    "upstream": [],
    "downstream": [
      "dag_31",
      "dag_21"
    ]
  },
  "dag_31": {
    "upstream": [
      "dag_11"
    ],
    "downstream": []
  }
}

DEPENDENCY_FILE = f'{os.path.dirname(os.path.abspath(__file__))}/dependencies.json'
def print_dependencies(dag_id):
    with open(DEPENDENCY_FILE, 'r+') as f:
        dependencies = json.load(f)
    upstream_dag_list = "\n".join(["- " + dag for dag in dependencies[dag_id]["upstream"]])
    downstream_dag_list = "\n".join(["- " + dag for dag in dependencies[dag_id]["downstream"]])
    return dedent('''\
    ### Dependencies
    #### Upstream DAGs
    {upstream_dag_list}
    #### Downstream DAGs
    {downstream_dag_list}
    '''
    ).format(upstream_dag_list=upstream_dag_list, downstream_dag_list=downstream_dag_list)
    
with DAG(
    'dag_11',
    ...
) as dag:
    ...
    dag.doc_md = print_dependencies(dag.dag_id)
    ...

import ast
...

DEPENDENCY_FILE = f'{os.path.dirname(os.path.abspath(__file__))}/dependencies.json'

class DependencyAnalyzer(ast.NodeVisitor):
    def __init__(self, filename):
        super().__init__()
        self.__dependencies = set()
        self.current_dag = ''
        self.filename = filename

    # DAGファイルのコードをパースして依存DAGをself.__dependenciesに設定する
    def analyze(self):
        with open(self.filename, 'r') as f:
            node = ast.parse(f.read())
        self.visit(node)

    # ExternalTaskSensorの呼び出しから引数external_dag_idの値を抽出する
    def visit_Call(self, node: ast.Call):
        func = node.func
        if isinstance(func, ast.Name):
            if func.id == 'ExternalTaskSensor':
                for kwarg in node.keywords:
                    if kwarg.arg == 'external_dag_id':
                        self.__dependencies.add(kwarg.value.value)
            elif func.id == 'DAG':
                self.current_dag = node.args[0].value

    @property
    def dependencies(self):
        return (self.current_dag, self.__dependencies)


def update_dependency(dag_id, dependencies):
    if os.path.exists(DEPENDENCY_FILE):
        with open(DEPENDENCY_FILE, 'r+') as f:
            dependencies_master_json = json.load(f)
    else:
        dependencies_master_json = {}

    # 対象のDAGのupstreamを設定・更新
    deps = dependencies_master_json.setdefault(dag_id, {'upstream': [], 'downstream': []})
    deps['upstream'] = list(dependencies)

    # upstreamのDAGのdownstreamに対象のDAGを設定
    for dependent_dag in dependencies:
        deps = dependencies_master_json.setdefault(dependent_dag, {'upstream': [], 'downstream': []})
        downstream = set(deps['downstream'])
        downstream.add(dag_id)
        deps['downstream'] = list(downstream)

    # 対象のDAGが依存していないDAGのdownstreamから対象のDAGを削除
    for dag, deps in dependencies_master_json.items():
        if dag not in dependencies:
            deps['downstream'] = list(set(deps['downstream']) - {dag_id})

    # dependencies.jsonに書き出す
    with open(DEPENDENCY_FILE, 'w+') as f:
        json.dump(dependencies_master_json, f)

if __name__ == '__main__':
    filename = sys.argv[1]  # DAGファイル名
    v = DependencyAnalyzer(filename)
    v.analyze()
    update_dependency(*v.dependencies)

# CI/CDでdependencies.jsonを更新
$ python dag_dependency.py dag_11.py