learn-sql-model/d3.py

70 lines
2.1 KiB
Python

import sqlite3
from jinja2 import Environment, FileSystemLoader
def get_tables_and_columns(conn):
cursor = conn.cursor()
cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
tables = [
{
"name": table[0],
"columns": get_columns(conn, table[0]),
"foreign_keys": get_foreign_keys(conn, table[0]),
}
for table in cursor.fetchall()
]
return tables
def get_columns(conn, table_name):
cursor = conn.cursor()
cursor.execute(f"PRAGMA table_info({table_name});")
columns = [row[1] for row in cursor.fetchall()]
return columns
def get_foreign_keys(conn, table_name):
cursor = conn.cursor()
cursor.execute(f"PRAGMA foreign_key_list({table_name});")
foreign_keys = [
{"id": row[0], "from": row[3], "to_table": row[2], "to": row[4]}
for row in cursor.fetchall()
]
return foreign_keys
def generate_links(tables):
links = []
for t_index, table in enumerate(tables):
for fk in table["foreign_keys"]:
target_index = next(
i for i, target in enumerate(tables) if target["name"] == fk["to_table"]
)
source_y = 40 + table["columns"].index(fk["from"]) * 20
target_y = 40 + tables[target_index]["columns"].index(fk["to"]) * 20
links.append(
{
"source": {"x": 50 + t_index * 150 + 120, "y": 50 + source_y},
"target": {"x": 50 + target_index * 150, "y": 50 + target_y},
}
)
return links
def generate_er_diagram(database_path):
conn = sqlite3.connect(database_path)
tables = get_tables_and_columns(conn)
links = [] # Currently, we won't extract relationships
links = generate_links(tables)
env = Environment(loader=FileSystemLoader("templates"))
template = env.get_template("er_diagram.html")
with open("index.html", "w") as f:
f.write(template.render(tables=tables, links=links))
if __name__ == "__main__":
db_path = "database.db"
generate_er_diagram(db_path)