Skip to content

Instantly share code, notes, and snippets.

@maresb
Last active September 8, 2021 21:15
Show Gist options
  • Save maresb/4c978e0e0ebf7379b3abc94bce2d050d to your computer and use it in GitHub Desktop.
Save maresb/4c978e0e0ebf7379b3abc94bce2d050d to your computer and use it in GitHub Desktop.
from typing import Dict, List, Union
from sqlalchemy import Column
from sqlalchemy.orm import (
ColumnProperty,
InstrumentedAttribute,
MapperProperty,
RelationshipProperty,
)
from sqlmodel.main import SQLModelMetaclass
COLUMNLIKE = Union[str, Column, ColumnProperty, InstrumentedAttribute]
def set_foreign_keys(
model: SQLModelMetaclass,
rels_with_fks: Dict[str, Union[COLUMNLIKE, List[COLUMNLIKE]]],
):
"""Written to solve https://github.com/tiangolo/sqlmodel/issues/10.
After declaring the SQLModel class, run something like:
set_foreign_keys(
Transaction,
{"to_account": "to_account_id", "from_account": "from_account_id"}),
)
"""
properties: Dict[str, MapperProperty] = model.__mapper__._props # type: ignore
columns: Dict[str, Column] = model.__table__.columns # type: ignore
for rel, fks in rels_with_fks.items():
if not isinstance(fks, List):
fks = [fks]
relationship_property = properties[rel]
assert isinstance(relationship_property, RelationshipProperty)
fk_column_instances = [_coerce_column(fk, columns) for fk in fks]
relationship_property._user_defined_foreign_keys = fk_column_instances # type: ignore # noqa
def _coerce_column(
col: COLUMNLIKE,
columns: Dict[str, Column],
) -> Union[Column, ColumnProperty]:
"""Convert COLUMNLIKE into a Column or ColumnProperty, if it isn't already."""
if isinstance(col, InstrumentedAttribute):
col = col.property
assert isinstance(col, ColumnProperty)
if isinstance(col, Column) or isinstance(col, ColumnProperty):
return col
if isinstance(col, str):
if "." in col:
raise ValueError(f"Please pass '{col}' as an object instead of a string.")
return columns[col]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment