Is it possible to load sqlalchemy relationship fields conditionally?

[citizen4371]

Hi everybody, I've got the following question:
Let's say I have an sqlalchemy class, a Post that has a one-to-one relation to Author, corresponding Pydantic schemas, and an endpoint:

# orm class
class OrmPost(Base):
    __tablename__ = 'posts'
    id = Column(Integer, primary_key=True, nullable=False)
    author_id = Column(String(50), ForeignKey('authors.author_id'))
    author = relationship('Author')

# pydantic schemas:
class Author(BaseModel):
   name: str

   class Config:
       orm_mode = True

class Post(BaseModel):
   id: int
   author_id: Optional[str]: None
   author: Optional[Author]: None

   class Config:
       orm_mode = True

# endpoint:
@app.get('/posts', response_model=List[Post])
def get_posts():
    with db_manager.session_scope() as session:
        return session.query(OrmPost).all()

Currently FastApi tries to load every field from the pydantic schema, which leads to author field in the example above to be lazy loaded for every post item. Is there a way to stop FastApi from loading relationship fields, so that the only scenario in which these fields are present in the response would be the case in which an explicit join is made prior to orm instance being processed by FastApi ? I would make the join conditionally, based on some query parameter:

@app.get('/posts', response_model=List[Post])
def get_posts(with_author=False):
    with db_manager.session_scope() as session:
        query = session.query(OrmPost)
        if with_author:
             query = query.options(joinedload(OrmPost.author))
        return query.all()

Environment

• OS: macOS:
• FastAPI Version: 0.61.1
• Python version: 3.7.5

[ArcLightSlavik]

Try this https://fastapi.tiangolo.com/tutorial/dependencies/

[citizen4371]

@ArcLightSlavik thanks, but I don't see how dependency injections could help in my case, could you elaborate, please ?

[citizen4371]

It seems that my issue is an inverse of #194, old relationship properties behavior described there is what I'm trying to achieve.

[citizen4371]

I think I found an ok solution using a custom GetterDict, here's what I came up with:

from typing import Any

from pydantic import BaseModel
from pydantic.utils import GetterDict
import sqlalchemy

class IgnoreLazyGetterDict(GetterDict):
    def __getitem__(self, key: str) -> Any:
        try:
            if self._is_lazy_loaded(key):
                return None

            return getattr(self._obj, key)
        except AttributeError as e:
            raise KeyError(key) from e

    def get(self, key: Any, default: Any = None) -> Any:
        # if a relationship field is not already loaded as a result of explicit join, ignore it,
        # so that pydantic schema doesn't trigger a lazy-load
        if self._is_lazy_loaded(key):
            return None

        return getattr(self._obj, key, default)

    def _is_lazy_loaded(self, key: Any) -> bool:
        return key in sqlalchemy.orm.attributes.instance_state(self._obj).unloaded

# a model to be used with sqlalchemy orm instances. It won't trigger lazy-load of relationship properties
class IgnoreLazyBaseModel(BaseModel):
    class Config:
        orm_mode = True
        getter_dict = IgnoreLazyGetterDict

[tiangolo]

Thanks for the help here @ArcLightSlavik ! ☕

I'm glad you found a solution @citizen4371 ! If that works for you, then you can close the issue. 🚀

In any case, you could also return Pydantic models directly. So, you could have 2 Pydantic models, one with the relationships and one without them. And then create the models in your code and return them.

Or you could also extract the data that you are sure you want to return in a dictionary or list of dictionaries, and return that.

[nickgieschen]

I think I found an ok solution using a custom GetterDict, here's what I came up with:

from typing import Any

from pydantic import BaseModel
from pydantic.utils import GetterDict
import sqlalchemy

class IgnoreLazyGetterDict(GetterDict):
    def __getitem__(self, key: str) -> Any:
        try:
            if self._is_lazy_loaded(key):
                return None

            return getattr(self._obj, key)
        except AttributeError as e:
            raise KeyError(key) from e

    def get(self, key: Any, default: Any = None) -> Any:
        # if a relationship field is not already loaded as a result of explicit join, ignore it,
        # so that pydantic schema doesn't trigger a lazy-load
        if self._is_lazy_loaded(key):
            return None

        return getattr(self._obj, key, default)

    def _is_lazy_loaded(self, key: Any) -> bool:
        return key in sqlalchemy.orm.attributes.instance_state(self._obj).unloaded

# a model to be used with sqlalchemy orm instances. It won't trigger lazy-load of relationship properties
class IgnoreLazyBaseModel(BaseModel):
    class Config:
        orm_mode = True
        getter_dict = IgnoreLazyGetterDict

I found that the above solution wasn't being applied to only relationships. This caused pydantic's orm_mode to fail. Below is the fix I came up with to use relationships only.

class IgnoreLazyGetterDict(GetterDict):
    def __getitem__(self, key: str) -> Any:
        try:
            if self._is_lazy_loaded(key):
                return None

            return getattr(self._obj, key)
        except AttributeError as e:
            raise KeyError(key) from e

    def get(self, key: Any, default: Any = None) -> Any:
        # if a relationship field is not already loaded as a result of explicit join, ignore it,
        # so that pydantic schema doesn't trigger a lazy-load
        if self._is_relationship(key) and self._is_lazy_loaded(key):
            return None

        return getattr(self._obj, key, default)

    def _is_lazy_loaded(self, key: Any) -> bool:
        return key in sqlalchemy.orm.attributes.instance_state(self._obj).unloaded

    def _is_relationship(self, key: Any):
        relationship_keys = [r.key for r in inspect(self._obj.__class__).relationships]
        return key in relationship_keys