set_foreign_keys.py

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]