bedrock.generators.generate-db-schema

 1"""
 2This script generates the SQL schema for the database.
 3
 4It works by introspecting the models in the `model/` directory and generating the SQL schema for them.
 5
 6It relies on SQLAlchemy's `CreateTable` function to generate the SQL schema.
 7"""
 8import re
 9import sys
10from enum import Enum
11
12from sqlalchemy.dialects import postgresql
13from sqlalchemy.schema import CreateTable, CreateIndex
14from bedrock.config.models import get_models
15
16
17def shorten_table_name(name: str) -> str:
18    """
19    Shorten a table name to an acronym.
20    :param name: The name of the table
21    """
22    return "".join([w[0] for w in f"{name}".split("_")])
23
24
25if __name__ == "__main__":
26    BASE_PATH = sys.argv[1]
27    sys.path.insert(0, f"{BASE_PATH}")
28    models = get_models(
29        f"{BASE_PATH}/model" if len(sys.argv) > 1 else "../model",
30        sys.argv[2] if len(sys.argv) > 2 else "model."
31    )
32
33    enum_statements = set()
34    for model in models:
35        for column in model.__table__.columns:
36            if issubclass(column.type.python_type, Enum) and column.type.enums:
37                values = ', '.join([f'\'{e}\'' for e in column.type.enums])
38                enum_statements.add(f"CREATE TYPE {column.type.name} AS ENUM ({values});")
39
40    create_statements = []
41    foreign_keys = []
42    indexes = []
43    reserved_words = ["name", "label", "message", "data", "sequence"]
44
45    comma_parenthesis_regex = re.compile(r",\s*(\n\s*)*\)")
46    referenced_table_regex = re.compile(r"REFERENCES (.*) \(")
47    for model in models:
48        create_statement = f"{CreateTable(model.__table__).compile(dialect=postgresql.dialect())}".replace("\n\n", "\n")
49        for word in reserved_words:
50            create_statement = re.sub(fr"(\s)({word}) (VARCHAR|INTEGER|JSONB)", r"""\1"\2" \3""", create_statement)
51        # Do Foreign Keys
52        matches = re.findall(r"(FOREIGN KEY.*)", create_statement)
53        if matches:
54            for m in matches:
55                referenced_table = referenced_table_regex.findall(m)[0]
56                fk_name = f"fk_{shorten_table_name(model.__table__)}_{shorten_table_name(referenced_table)}"
57                fk_statement = f"ALTER TABLE {model.__table__}  ADD CONSTRAINT {fk_name} {m};".replace(", ;", ";")
58                foreign_keys.append(fk_statement)
59                create_statement = comma_parenthesis_regex.sub("\n)", create_statement.replace(m, ""))
60        create_statements.append(f"{create_statement};".replace("\n;", ";"))
61        # Do Indexes
62        indexes.extend(f"{CreateIndex(idx).compile(dialect=postgresql.dialect())};" for idx in model.__table__.indexes)
63
64    print("\n".join(enum_statements))
65    print("\n")
66    print("\n".join(create_statements))
67    print("\n")
68    print("\n".join(foreign_keys))
69    print("\n")
70    print("\n".join(indexes))
71    print("\n")