fastapi / sqlmodel

SQL databases in Python, designed for simplicity, compatibility, and robustness.
https://sqlmodel.tiangolo.com/
MIT License
14.23k stars 630 forks source link

make relationship nicely inheritable #167

Open 5cat opened 2 years ago

5cat commented 2 years ago

First Check

Commit to Help

Example Code

from typing import Optional

from sqlmodel import Field, Session, SQLModel, create_engine, select
from sqlalchemy.orm import declared_attr, relationship

class Team(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str
    headquarters: str

class HeroBase(SQLModel):
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str

    team_id: Optional[int] = Field(default=None, foreign_key="team.id")

    @declared_attr
    def team(self) -> Optional[Team]:
        return relationship("Team")  # type: ignore

class AnimalHero(HeroBase, table=True):
    species: str

class RobotHero(HeroBase, table=True):
    model_number: int

engine = create_engine("sqlite://")

SQLModel.metadata.create_all(engine)

with Session(engine) as session:
    peace = Team(name="peace", headquarters="earth")
    evil = Team(name="evil", headquarters="mars")
    session.add_all((peace, evil))
    session.commit()

    chickenman = AnimalHero(name="chickenman", species="chickens", team_id=peace.id)
    siri = RobotHero(name="siri", model_number=0x6af, team_id=evil.id)
    session.add_all((chickenman, siri))
    session.commit()

with Session(engine) as session:
    animals = session.exec(select(AnimalHero)).all()
    robots = session.exec(select(RobotHero)).all()
    print(f"{animals=}")
    print(f"{robots=}")
    assert all(hasattr(h, "team") for h in animals)
    assert all(hasattr(h, "team") for h in robots)
    assert animals == [AnimalHero(id=1, species='chickens', team_id=1, name='chickenman')]
    assert robots == [RobotHero(id=1, model_number=1711, team_id=2, name='siri')]

    print(f"{[h.team for h in animals]=}")
    print(f"{[h.team for h in robots]=}")
    assert [h.team for h in animals] == [Team(name='peace', id=1, headquarters='earth')]
    assert [h.team for h in robots] == [Team(name='evil', id=2, headquarters='mars')]

Description

I want to create an SQLModel base class that contains a relationship which i can inherit. in the example there is a HeroBase which is inherited by AnimalHero and RobotHero. It inherits the columns correctly but with the relationships, i need to use sqlalchemy.orm.declared_attr, sqlalchemy.orm.relationship and # type: ignore so the type checker doesnt get mad.

then in insertion i need to first insert the teams, commit, then link the heroes via the team_id manually instead of doing AnimalHero(name="chickenman", species="chickens", team=peace), so the current way is very similar to this

Wanted Solution

I just want to inherit the relationship provided by SQLModel.Relationship and doing the insertion this way.

Wanted Code

from typing import Optional

from sqlmodel import Field, Session, SQLModel, create_engine, select, Relationship

class Team(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str
    headquarters: str

class HeroBase(SQLModel):
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str

    team_id: Optional[int] = Field(default=None, foreign_key="team.id")

    team: Team = Relationship()

class AnimalHero(HeroBase, table=True):
    species: str

class RobotHero(HeroBase, table=True):
    model_number: int

engine = create_engine("sqlite://")

SQLModel.metadata.create_all(engine)

with Session(engine) as session:
    peace = Team(name="peace", headquarters="earth")
    evil = Team(name="evil", headquarters="mars")
    chickenman = AnimalHero(name="chickenman", species="chickens", team=peace)
    siri = RobotHero(name="siri", model_number=0x6af, team=evil)
    session.add_all((chickenman, siri))
    session.commit()

with Session(engine) as session:
    animals = session.exec(select(AnimalHero)).all()
    robots = session.exec(select(RobotHero)).all()
    print(f"{animals=}")
    print(f"{robots=}")
    assert all(hasattr(h, "team") for h in animals)
    assert all(hasattr(h, "team") for h in robots)
    assert animals == [AnimalHero(id=1, species='chickens', team_id=1, name='chickenman')]
    assert robots == [RobotHero(id=1, model_number=1711, team_id=2, name='siri')]

    print(f"{[h.team for h in animals]=}")
    print(f"{[h.team for h in robots]=}")
    assert [h.team for h in animals] == [Team(name='peace', id=1, headquarters='earth')]
    assert [h.team for h in robots] == [Team(name='evil', id=2, headquarters='mars')]

Alternatives

Using directly SQLAlchemy instead of SQLModel

Operating System

Linux

Operating System Details

No response

SQLModel Version

0.0.4

Python Version

Python 3.9.8

Additional Context

The reason for me to go to this route is to implement generic tables/associations and tried to look for examples in SQLAlchemy and tried to use SQLAlchemy-Utils generic_relationship with Relationship(sa_relationship=generic_relationship("object_id", "object_type")), although the later works for insertion, it doesnt work when you try to getattr the relationship after selection.

The wanted code is reasonable considering this is my second day using SQLModel and it feels intuitive to do it that way.

kellen commented 2 years ago

+1 for this working roughly the same way as @declared_attr does in SQLAlchemy: https://docs.sqlalchemy.org/en/13/orm/extensions/declarative/mixins.html#mixing-in-columns-in-inheritance-scenarios

Currently, the FK fields generate the correct DB columns, but the relationship attrs are dropped silently.

thomasborgen commented 2 years ago

+1 - I couldn't figure out why Relationships were not working when I had declared them in a "base" class.

The reason why I put it in the Base class was that I was sure it would work the same as with Fields like written here: https://sqlmodel.tiangolo.com/tutorial/fastapi/multiple-models/?h=base#multiple-models-with-inheritance and https://sqlmodel.tiangolo.com/tutorial/fastapi/multiple-models/?h=base#multiple-models-with-inheritance

It would be awesome if Relationships could work the same.