I am looping through my SQLModel classes, extracting the field info, and creating a .dbml format file from the details. This works great, except trying to find the foreign_key information in the fields. Here is an example of an SQLModel classes:
class Account(PostgresBase, table=True):
__tablename__ = 'accounts'
id: str = Field(default_factory=lambda: genb58_id(22, "account"),
primary_key=True, index=True, nullable=False)
first_name: str = Field(index=True)
last_name: str = Field(index=True)
username: str = Field(index=True, unique=True)
# ----- Relationships -----
some_activity_records: list["SomeActivity"] = Relationship(
back_populates="account")
class SomeActivity(PostgresBase, table=True):
__tablename__ = 'some_activity_recs'
id: str = Field(default_factory=lambda: genb58_id(22, "activity"),
primary_key=True, index=True, nullable=False)
name: str
type: Union[dict, None] = Field(sa_type=JSON)
fk_account_id: Union[str, None] = Field(
default=None, foreign_key="accounts.id")
things: list[str] = Field(sa_type=ARRAY(String), default=[])
# ----- Relationships -----
account: Union["Account", None] = Relationship(
back_populates="some_activity_records")
How can I tease out the foreign_key info that the SomeActivity fk_account_id field is linked to the table+field accounts.id?
Here is the script I am running to get the fields and the relationships:
import inspect
from sqlmodel import SQLModel
import sys
import types
import typing
from typing import get_args
from models import *
# Get all classes from models.py that are subclasses of SQLModel.
models = [obj for name, obj in inspect.getmembers(sys.modules[__name__])
if inspect.isclass(obj) and issubclass(obj, SQLModel) and obj != SQLModel]
# Generate dbml for each model.
dbml_content = ""
for model in models:
# Extract class docstring.
doc = inspect.getdoc(model).replace('\n', ' ')
if doc.startswith("Usage docs"):
doc = None
# Build the table definition.
dbml_content = f"Table {model.__tablename__} {{\n"
dbml_content += f" // {doc}\n" if doc else ""
# Build the field lines.
for key in model.model_fields.keys():
fieldstr = f" {key}"
field_info = model.model_fields[key]
annotation = field_info.annotation
match type(annotation):
case t if t == type:
type_name = annotation.__name__
if type_name == "str":
type_name = "varchar"
opts = ""
if key == "id":
opts = "[pk, unique, not null]"
fieldstr += f" {type_name} {opts}"
case t if t == typing._UnionGenericAlias:
args = get_args(annotation)
arg_name = args[0].__name__
match arg_name:
case "dict":
arg_name = "json"
case "str":
arg_name = "varchar"
fieldstr += f" {arg_name}"
case t if t == types.GenericAlias:
fieldstr += f" varchar[]"
case t if t == typing._LiteralGenericAlias:
args = get_args(annotation)
fieldstr += f" varchar //{args}"
dbml_content += fieldstr + "\n"
dbml_content += "}\n\n"
# Add relationships.
relationships = []
for model in models:
for field in model.model_fields.values():
if 'foreign_key' in field.metadata:
fk = field.metadata['foreign_key']
if fk:
ref_table, ref_field = fk.split('.')
relationships.append(
f'Ref: "{model.__tablename__}"."{field.name}" < "{ref_table}"."{ref_field}"\n')
dbml_content += ''.join(relationships)
print(dbml_content)
The "Generate dbml for each model." loop works great and creates an appropriate dbml, eg:
Table accounts {
id varchar [pk, unique, not null]
first_name varchar
last_name varchar
username varchar
}
Table some_activity_recs {
id varchar [pk, unique, not null]
name varchar
type json
fk_account_id varchar
things varchar[]
}
However, when I run the section to "Add relationships.", the field.metadata is always an empty list, even when inspecting the fk_account_id field.
I have tried to use the field.foreign_key member, which should exist in SQLModel FieldInfo class, but it gives me the following traceback:
***** model.__name__: Account
***** type(model): <class 'sqlmodel.main.SQLModelMetaclass'>
***** field.foreign_key: PydanticUndefined
Traceback (most recent call last):
File "/soc-api/./soc/dao/scripts/create_dbml.py", line 86, in <module>
print(f"***** field.foreign_key: {field.foreign_key}")
^^^^^^^^^^^^^^^^^
AttributeError: 'FieldInfo' object has no attribute 'foreign_key'
The problem I ran into has to do with the fact that some fields in the models loop are <class sqlmodel.main.FieldInfo> while others happen to be <pydantic.fields.FieldInfo> even though the model type always came back as <class 'sqlmodel.main.SQLModelMetaclass'>, as described above.
The answer is to just use this conditional in the loop:
for model in models:
for key in model.model_fields.keys():
field_info = model.model_fields[key]
if hasattr(field_info, "foreign_key") and isinstance(field_info.foreign_key, str):
print(f"***** table_name: {model.__tablename__}")
print(f"***** field_name: {key}")
print(f"***** field_info.foreign_key: {field_info.foreign_key}")
The output is exactly what I need, the table and field pointed to by the foreign key:
***** table_name: some_activity_recs
***** field_name: fk_account_id
***** field_info.foreign_key: accounts.id