[Python] Use lazy/eager loading dynamically with SQLModel

I have one issue with SQLModel atm which I cant figure out.

The Relationship()'s are super handy and make the code clean. By default sqlalchemy uses lazy loading. The issue is the children in the relationship return a "greenlet error" rather than the child objects.

As stated in this StackOverflow answer: "The main reason for this error is that in the synchronous sqlalchemy driver you can use the session for the lazy loading queries properly. Assume you have two relational models to each other, fetch the parent by a query then you are going to reach the child leading to another query (i.e. lazy loading). However, when you are using the async driver on top of sqlalchemy (e.g. asyncpg) the session will be automatically closed each time after a query leading to the raising error for getting the child info. "

The solution is to pass "lazy": "selectin" or some other eager loading strategy when you define the Relationship(). The issue then is that not every time I use the model I want to use eager loading. Sometimes I want to use the Conversation model without getting all the messages. Is there a way to set this when I call select? I tried sqlalchemy .options but this doesn't seem to pass through SQLModel?

from typing import Optional
from uuid import UUID, uuid4

from sqlmodel import SQLModel, Relationship, Field


class Citation(SQLModel, table=True):

id: Optional[UUID] = Field(
default_factory=uuid4,
primary_key=True,
description="The unique identifier of the citation",
)
content: str
message_id: UUID | None = Field(
foreign_key="message.id",
ondelete="CASCADE",
index=True,
description="The unique identifier of the message the citation belongs to",
)
message: "Message" = Relationship(
back_populates="citations", sa_relationship_kwargs={"lazy": "selectin"}
)


class Message(SQLModel, table=True):

id: Optional[UUID] | None = Field(
default_factory=uuid4,
primary_key=True,
description="The unique identifier of the message",
)
conversation_id: UUID | None = Field(
foreign_key="conversation.id",
ondelete="CASCADE",
index=True,
description="The unique identifier of the conversation the message belongs to",
)
content: str
conversation: "Conversation" = Relationship(
back_populates="messages", sa_relationship_kwargs={"lazy": "selectin"}
)
citations: list["Citation"] = Relationship(
back_populates="message",
sa_relationship_kwargs={"lazy": "selectin", "cascade": "all, delete-orphan"},
)


class Conversation(SQLModel, table=True):

id: Optional[UUID] | None = Field(
default_factory=uuid4,
primary_key=True,
description="The unique identifier of the conversation",
)
title: str
creator_id: int
messages: list["Message"] = Relationship(
back_populates="conversation",
sa_relationship_kwargs={"lazy": "selectin", "cascade": "all, delete-orphan"},
)

# Example of my query method
async def get_conversations(
session: AsyncSession, chatbot_id: UUID
) -> Sequence[Conversation]:
"""Get all conversations for the chatbot"""
conversations = (
await session.exec(
select(Conversation).where(Conversation.creator_id == 1)
)
).all()
return conversations

Continue reading...