Files
workout/scripts/generate_schema.py
2026-01-30 18:47:26 +11:00

81 lines
2.3 KiB
Python

import os
import sys
# Add the project root to sys.path so we can import 'features'
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import subprocess
import tempfile
import psycopg2
from psycopg2.extras import RealDictCursor
from urllib.parse import urlparse
from dotenv import load_dotenv
# Import Schema logic
from features.schema import Schema
def main():
# Load environment variables from .env
load_dotenv()
database_url = os.environ.get('DATABASE_URL')
if not database_url:
print("Error: DATABASE_URL not set in environment.")
return
# Connect to the database
db_url = urlparse(database_url)
conn = psycopg2.connect(
database=db_url.path[1:],
user=db_url.username,
password=db_url.password,
host=db_url.hostname,
port=db_url.port
)
def execute_func(query, args=()):
cur = conn.cursor(cursor_factory=RealDictCursor)
cur.execute(query, args)
rv = None
if cur.description is not None:
rv = cur.fetchall()
cur.close()
return rv
# Initialize Schema logic
schema_logic = Schema(execute_func)
print("Fetching schema information...")
schema_info = schema_logic.get_schema_info()
print("Generating Mermaid code...")
mermaid_code = schema_logic.generate_mermaid_er(schema_info)
print("Rendering SVG via mermaid-cli...")
try:
with tempfile.TemporaryDirectory() as tmpdir:
input_file = os.path.join(tmpdir, "input.mmd")
output_file = os.path.join(tmpdir, "schema.svg")
# Destination path
target_file = os.path.abspath(os.path.join("static", "img", "schema.svg"))
os.makedirs(os.path.dirname(target_file), exist_ok=True)
with open(input_file, "w") as f:
f.write(mermaid_code)
# Run mmdc
subprocess.run(
["bun", "x", "mmdc", "-i", input_file, "-o", target_file, "-b", "transparent"],
check=True
)
print(f"Successfully generated {target_file}")
except Exception as e:
print(f"Error during SVG generation: {e}")
finally:
conn.close()
if __name__ == "__main__":
main()