SQLスクリプト上でのテーブルや共通テーブル式の依存関係を可視化する

電通デジタルでデータサイエンティストを務める吉田です。
こちらはDentsu Digital Tech Blogアドベントカレンダー16日目の記事です。

今回はSQLスクリプトを解析して参照されているテーブルや共通テーブル式の依存関係を可視化するツールを作成したのでご紹介いたします。

どのようなものを作ったか

電通デジタルデータサイエンスグループでは普段、主にBigQuery上でSQLスクリプトを記述してデータの抽出や集計・整形作業を行うことが多いです。

最近ではBigQuery MLでBoostingTreeやAutoML Tablesのモデルが作成できるようになるなど機能が拡張されており、BigQuery上でデータを前処理してBigQuery MLで機械学習モデルを生成、といったようにデータ管理からモデリングまでBigQuery上で完結するシーンもみられるようになってきました。

シンプルな構文でデータテーブルへの問い合わせができるSQLですが、手の込んだ集計・データ整形をしようとすると非常に長く煩雑なスクリプトになることもしばしばあります。

データの集計項目が多い時や、機械学習用のデータ作成でラベル変換や特徴量生成をSQLで行うときなどがそれにあたり、時には数百行ほどの長さに至ります。

煩雑なデータ整形をする際には、深すぎる副問い合わせを避けてスクリプトの可読性を維持するために共通テーブル式（WITH句）を使うことが多いと思いますが、これでも書いてからしばらく時間が経った長文SQLを再度読み解こうとすると、参照しているテーブルやテーブル式を探してスクリプトを上下に行ったり来たり...と多少骨が折れます。

そこで、ソフトウェア設計の構造概要を示すUML図のような形でテーブル・テーブル式の参照・被参照関係を可視化できるものを作りました。
SQLスクリプトを入力すると下記のような図を表示するようなものになります。

実装

今回はPythonを使い、main.pyファイルに実装していきます。
Webアプリのように使いたかったので、PythonでシンプルなWebフロントエンド構築が容易にできるStreamlitを使用しました。

Streamlit — The fastest way to create data apps

下記コマンドでStreamlitをインストールします。

pip install streamlit

Streamlitではテキストフォームに文字列を入力し、Ctrl+Enterで入力した文字列に対して指定した処理を実行することができます。
フォームにSQLスクリプトを入力し、図を表示できるようにします。

import streamlit as st

# text form
query_script = st.text_area('Enter SQL script.', height=400)

try:
    # if SQL script has been entered, parse that and display the figure
    ...
except:
    # if SQL syntax is wrong or have not been entered raise message
    st.subheader('Check and modify script')
    

スクリプトの構文解析にはregexモジュール（標準のreモジュールの拡張版）を使用しています。

最初に--や#、/*で始まるようなコメント行をすべて削除しておきます。
with hoge as (...や, fuga as (...とマッチするような部分文字列を正規表現で探索し、テーブル式名とクエリを取得します。

WITH句があった場合となかった場合でパターン分けをして主問い合わせ部分の文字列も取得します。

import regex
# delete comment
query_script = regex.sub(r'--.*\n|#.*\n|/\*([^*]|\*[^/])*\*/', '', query_script)

# get CTEs
cte_pattern = r'(?:with|,)\s*(\w+)\s+as\s*(?<rec>\((?:[^\(\)]+|(?&rec))*\))'
ctes = regex.finditer(cte_pattern, query_script, regex.IGNORECASE)
queries = {cte.group(1): cte.group(2) for cte in ctes}

# get main query
main_pattern = r'\)[;\s]*select' if any(queries) else r'select'
start_main = regex.search(main_pattern, query_script, regex.IGNORECASE).span()[0] + 1
queries['main'] = query_script[start_main:].strip()

# find reference table or CTEs
ref_pattern = r'(?:from|join)\s+([`.\-\w]+)'
dependencies = dict()
for name, script in queries.items():
   refs = regex.findall(ref_pattern, script, regex.IGNORECASE)
   dependencies[name] = [ref for ref in refs]

作成した辞書を元にテーブルとテーブル式の依存関係を作図します。
今回はdot言語でグラフを描画できるGraphvizを使用しました。

PythonからGraphvizを使用する際には、Python用のクライアントライブラリのインストールと併せてGraphviz本体のインストールが必要になります。

ライブラリの方は下記コマンドでインストール可能です。

pip install graphviz

Graphviz本体は公式サイトからダウンロードしてインストールすることも可能ですが、Windows環境でscoopなどのパッケージ管理環境を構築していれば、下記コマンドでインストール可能です。

scoop install graphviz 

辞書を読み込み、参照元から参照先に矢印が向かうようにノードとエッジを生成していきます。主問い合わせ部（main）とテーブル、共通テーブル式で色分けして表示できるように描画オプションを指定します。

st.graphviz_chart()を用いて生成したGraphvizオブジェクトを図としてStreamlit上で表示します。

from graphviz import Digraph

# draw graph
g = Digraph()
g.attr('node', shape='box')
with g.subgraph(name='querys') as c:
   c.attr(color='blue', label='querys')
   for name, refs in dependencies.items():
       for ref in refs:
           c.node(name, style='bold, filled' if name=='main' else 'solid, filled', fillcolor='#FFFFFF' if name=='main' else '#80CBC4')
           c.node(ref, style='solid, filled', fillcolor='#81C784' if '.' in ref else '#80CBC4')
           c.edge(name, ref)

st.graphviz_chart(g)

実行

Streamlitを起動して実際に動かしてみます。

streamlit run main.py

ブラウザが起動し下記のようなシンプルなテキストフォームが表示されます。

試しに適当なSQLスクリプトを入力してみると、依存関係通りにグラフが生成・出力されることがわかります。

実際に現場で作成して使っている長大SQLスクリプトを入力してみると以下のような出力になりました。
かなり複雑なクエリになっていますが、どのテーブル式がどこを参照しているのかが把握しやすくなりました。※テーブル名は伏せております。

Streamlitはデータ可視化などを中心にデータ分析業務をサポートする形での使い方を想定して開発されているようですが、それ以外のシーンでもごく簡単な社内ツールを作る際などに非常に取り入れやすいなと感じました。
また、今回実装したものを応用してよりデータエンジニアリング作業を進めやすいスクリプトエディタUIを作れないかとも模索し始めています。

次回、アドベントカレンダー17日目の記事は「Micro Frontends導入の覚書」です。

参考

・Streamlit - The fastest way to build and share data apps
・Graphviz document