dbtで見やすいER図を生成する

背景: dbtを使っていてもER図は欲しい!

dbtはモデル間のリネージなど可視化が得意なツールではありますが、万能なわけではありません。モデルの生成過程などはリネージで担保できますが、分析時に「どれとどのモデルがJOINできて、JOINする際のキーはこれを使って」というER図で扱うような可視化はディフォルトではできません。

DWHを作っている側からすると「このテーブルはあの辺のテーブルと一緒に使うと便利で、いつもあのキーでJOINして」というのが頭の中に入っていることが多いため、ER図がなくてもどうにかなることも多いでしょう。しかし、分析に慣れていない人や分析に慣れている人であっても、普段と異なるドメインのテーブルを触るときはER図が提供してくれる情報は有用です。ちなみに私は鳥頭なので、ER図は常に欲しいです。。

どうやってER図を生成するか

BigQueryのようなDWHはプロダクト側で使っているRDBと違って、最初から外部キーの設計などがされていることは稀でしょう。NoSQLを使っている場合、そういった情報をDBに組み込みにくい、といった事情があるかもしれません。何かしら後付けで外部キーなどのメタ情報を付与して、ER図相当の情報を構成していくことになると思います。

何かしらのメタデータを付与していくとして、ただのメタデータで終わってしまうとちょっともったいない感じがします(貧乏性)。dbtの文脈でいうとrelationships(参照整合性を満たしているかどうか)というGeneric testがあり、これが今回のユースケースに向いているように思います。このrelationshipsの引数の情報をうまく料理すれば、ER図のエッジの情報が生成できそうです。品質を高めるテストがER図の生成にも役立てば一石二鳥ですね。dbtでのrelationshipsのテストは以下のような形で書くことが多いです。

version: 2

models:
  - name: orders
    columns:
      - name: customer_id
        tests:
          - relationships:
            to: ref('customers')
            field: id
            config:
              # 運悪く参照整合性が満たされていない場合でも
              # ER図を生成したい場合はテストが必ず通るように調整しておく
              # もちろん通る場合は不要
              where: "false"

実際、世の中にはこういった後付けの情報を加工して、ER図を生成するツールを公開してくれている人がいます。ER図の出力形式も複数サポートされていていいですね。READMEを読みながら試してみると、割とすぐ動きました。

なお、relationが1:Nかといった情報はmetaに付与することができるので、よりリッチで正確な情報を付与したい場合はやってみるとよいでしょう。

どうやってER図を見やすくするか

さて、すぐ動くのはよいのですが、実用するにはもう一歩考える必要がありました。実際に動かしてみると、以下の問題がありました。

  • dbtのモデル配下の全てを対象にすると、ER図が巨大になり過ぎてもはや何も分からない
    • タグなどでdbtのselectorで対象を絞ることもできますが、自社の場合だとそれでも数十以上が対象となることが多く、これでもまだ実用に耐えません
  • 可視化する対象を手動で指定すればもちろん人間の目でも分かる大きさ になりますが、手動で指定して回るのは案外面倒です

そこで、ER図の可視化の対象を以下のように絞ることにしました。

  • 1: 対象となるテーブルを一個決める
  • 2: そのテーブルにrelationshipsで紐付いているテーブルからNホップでJOINで紐付けられる範囲を可視化の範囲とする

1のテーブルを変数と見做して可視化の出力を回せば、dbtのパッケージ内全てを一応可視化できていることにはなります。2でNホップの繋がりのテーブルに限定すれば、人間が見て分かる形になります。

これを実現するスクリプトを書きました。こういう感じで動かすと、モデルファイルと同じディレクトリにER図の情報を含むmarkdownファイルが生成されます。

% python scripts/expand_select_nodes_for_dbtedr.py \
  --manifest_filename target/manifest.json \
  --target_node_name model.my_model.hub_user \
  --max_hops 3 \
  --max_nodes 10

ER図をいい感じに生成するスクリプト(クリックで開きます)

import argparse
import json
import os
import subprocess
import sys
from itertools import chain
from typing import Any, Dict, List, Set


# ER図を出力するのに必要なrelationshipsを含むノードの情報を抽出する関数
def extract_relationships_nodes(manifest: Dict[str, Any]) -> List[Dict[str, Any]]:
    result = []
    for node in manifest["nodes"].values():
        if (
            node["resource_type"] == "test"
            and node.get("test_metadata", {}).get("name") == "relationships"
        ):
            result.append(node)
    return result


# target_nodesをスタートとして、Nホップ先のノードを返す関数
def search_n_hops_nodes(
    relationships_nodes: List[Dict[str, Any]],
    target_nodes: Set[str],
    max_hops: int,
    max_nodes: int = 10,  # ホップ数が多い場合、ER図に登場するノード数が多いため、上限を設ける
    n: int = 0,
) -> Set[str]:
    if n >= max_hops:
        return target_nodes

    initial_target_nodes = target_nodes.copy()
    for target_node_name in initial_target_nodes:
        for node in relationships_nodes:
            if target_node_name in node["depends_on"]["nodes"]:
                for new_node in node["depends_on"]["nodes"]:
                    if len(target_nodes) >= max_nodes:
                        print(
                            f"ER図内にあまりに多くのノード(n={len(target_nodes)})が含まれるため、出力を制限します",
                            file=sys.stderr,
                        )
                        return target_nodes
                    target_nodes.add(new_node)
    if initial_target_nodes == target_nodes:
        return target_nodes
    return search_n_hops_nodes(
        relationships_nodes,
        target_nodes,
        max_hops=max_hops,
        max_nodes=max_nodes,
        n=n + 1,
    )


def run_dbterd(n_hops_nodes: Set[str]) -> None:
    dbterd_command = [
        "dbterd",
        "run",
        "--target",
        "mermaid",
        "--entity-name-format",
        "schema.table",
        "--resource-type",
        "source",
        "--resource-type",
        "model",
    ]
    dbterd_select_targets = list(
        chain.from_iterable([["--select", f"exact:{node}"] for node in n_hops_nodes])
    )
    subprocess.run(dbterd_command + dbterd_select_targets)


def get_markdown_file_path(manifest: Dict[str, Any], node_name: str) -> str:
    filename_without_extension, _ = os.path.splitext(
        manifest["nodes"][node_name]["original_file_path"]
    )
    return filename_without_extension + ".md"


# markdown内のリンクがうまく動作するように相対パスを取得する関数
def get_relative_path(base_file_path: str, target_file_path: str) -> str:
    base_dir = os.path.dirname(base_file_path)
    return os.path.relpath(target_file_path, base_dir)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--manifest_filename",
        type=str,
        required=True,
    )
    parser.add_argument(
        "--target_node_name",
        type=str,
        required=True,
    )
    parser.add_argument(
        "--max_hops",
        type=int,
        required=True,
    )
    parser.add_argument(
        "--max_nodes",
        type=int,
        required=True,
    )
    args, _ = parser.parse_known_args()
    target_node_name = args.target_node_name

    with open(args.manifest_filename, "r") as file:
        manifest = json.load(file)

    n_hops_nodes = search_n_hops_nodes(
        extract_relationships_nodes(manifest),
        {target_node_name},
        max_hops=args.max_hops,
        max_nodes=args.max_nodes,
    )

    if len(n_hops_nodes) <= 1:
        print(
            "ER図を出力するのに十分なノード数がありません。ER図を書きたい場合はrelationshipsのテストを追加することを検討してください",
            file=sys.stderr,
        )
        sys.exit(1)

    run_dbterd(n_hops_nodes)

    with open("target/output.md", "r") as file:
        mermaid = file.read()

    markdown_file_path = get_markdown_file_path(manifest, target_node_name)
    with open(f"{markdown_file_path}", "w") as file:
        file.write(f"## {target_node_name}のER図\n")
        file.write("```mermaid\n")
        file.write(mermaid)
        file.write("```\n")
        file.write("## 関連するモデルのER図へのリンク\n")
        for node in n_hops_nodes:
            if node != args.target_node_name:
                node_file_path = get_relative_path(
                    markdown_file_path, get_markdown_file_path(manifest, node)
                )
                file.write(f"- [{node}]({node_file_path})\n")
    print(f"{markdown_file_path}にER図を出力しました")


if __name__ == "__main__":
    main()

実際に生成されたER図はこういう感じです。markdownで出力するため、関連するER図にさっと飛べますし、可視化の範囲がJOINできる範囲で制限されているので見やすいかと思います。

1つのモデルに対して上記のようなER図が生成されるので、全てのモデルに対してER図が最新のものに更新されていくようにGitHub ActionsなどでCIを週次で回して運用していくとよさそうです。

まとめ

  • ER図はdbtを使う上でももちろん有用です
  • dbtに付属のrelationshipsの参照照合性のテストを記述する形でER図の生成に必要な情報を付与することができます
  • dbtのリポジトリ全体のER図を可視化すると巨大過ぎるため、範囲をいい感じに絞るスクリプトを作りました

というわけで、ER図を頼りにDWHの開発者もDWHの利用者もHappyになっていきましょう。