クエリからリレーションシップを推定するPythonコード

はじめに

データカタログをどう管理すべきか考えていたところ、下記の記事に辿りついた。
データカタログを作成してZOZOTOWNデータベース定義をまとめた話

ZOZOTOWNでは、外部キー制約は貼られていないがリレーションシップを持っているものとして
運用されているテーブルが多数存在しており、その関係を手動でまとめるのは困難でした。

記事によると、共感する課題感があった。

本番環境で実際に実行されたクエリを大量に入手。
その中からJOIN句でつながっているテーブル名とカラム名を抜き出し、
どちらかがPKであれば仮想外部キーの候補とする。

と記載があったが、肝心の抜粋方法の記載はなかったので、
ChatGPTとともにクエリからリレーションシップ推定するPythonコードを作成する。

やりたいこと

クエリからリレーションシップ推定するPythonコードを作成する。
今回はリレーションがありそうってことを推定するとこまでを実施する。
1:N or 1:1の推定の問題は一旦先送りにする。

やったこと

• サンプルクエリ(Bigqueryっぽい)

SELECT 
    employees.name
    , departments.department_name
FROM `prd-pjt.employees` employees
INNER JOIN `prd-pjt.departments` departments
    ON employees.department_id = departments.department_id
LEFT JOIN `prd-pjt.salaries` salaries
    ON employees.employee_id = salaries.employee_id
WHERE employees.salary > 50000

• Pythonコード

import pandas as pd

sql_query = """
SELECT 
    employees.name
    , departments.department_name
FROM `prd-pjt.employees` employees
INNER JOIN `prd-pjt.departments` departments
    ON employees.department_id = departments.department_id
LEFT JOIN `prd-pjt.salaries` salaries
    ON employees.employee_id = salaries.employee_id
WHERE employees.salary > 50000
"""

# SQL文を一単語ずつに分割
words = sql_query.split()

# 空のデータフレームを作成
df = pd.DataFrame(columns=[
	"table_name","key_name","join_key_name","join_table_name"])

j = -1
syokai =  0
# 単語を順番にチェック
for i, word in enumerate(words):
    if word == "FROM" or word == "from":
        j = j + 1
        syokai =  0
        df.at[j, "table_name"] = words[i+1]
    elif word == "JOIN" or word == "join":
        if syokai == 0:
            df.at[j, "join_table_name"] = words[i+1]
        else :
            j = j + 1
            df.at[j, "table_name"] = df.at[j-1, "table_name"]
            df.at[j, "join_table_name"] = words[i+1]
        syokai = syokai + 1
    elif word == "ON" or word == "on":
        if '.' in words[i+1]:
            df.at[j, "key_name"] = words[i+1].split('.')[1]
        else:
            df.at[j, "key_name"] = words[i+1]
        if '.' in words[i+3]:
            df.at[j, "join_key_name"] = words[i+3].split('.')[1]
        else:
            df.at[j, "join_key_name"] = words[i+3]

df = df[df.notna().all(axis=1)]
df = df[df['table_name'].str.contains('prd-pjt')]
df = df[df['join_table_name'].str.contains('prd-pjt')]
df['key_name'] = df['key_name'].str.replace("=", "")
df['join_key_name'] = df['join_key_name'].str.replace("=", "")
df['table_name'].sort_values(ascending=False)
df = df.drop_duplicates()
df.head()

• 出力結果

	table_name	key_name	join_key_name	join_table_name
0	prd-pjt.employees	department_id	department_id	prd-pjt.departments
1	prd-pjt.employees	employee_id	employee_id	prd-pjt.salaries

課題

やり残したことは、以下の２点

• 無理やり単語を一個ずつ判定したんで、もっといい方法があるはず。
• 1:N or 1:1の推定の問題