How to query a nested attribute in SQLModel?

1.9k Views Asked by At

I have these two tables named User and UserRole.

import enum
from typing import List
from sqlmodel import Column, Enum, Field, Relationship, SQLModel


class UserRoleType(str, enum.Enum):
    admin = 'admin'
    client = 'client'


class UserRole(SQLModel, table=True):
    __tablename__ = 'user_role'
    id: int | None = Field(default=None, primary_key=True)

    type: UserRoleType = Field(
        default=UserRoleType.client,
        sa_column=Column(Enum(UserRoleType)),
    )

    write_access: bool = Field(default=False)
    read_access: bool = Field(default=False)

    users: List['User'] = Relationship(back_populates='user_role')


class User(SQLModel, table=True):
    id: int | None = Field(default=None, primary_key=True)
    username: str = Field(..., index=True)

    user_role_id: int = Field(..., foreign_key='user_role.id')
    user_role: 'UserRole' = Relationship(back_populates='users')

I can easily insert entries into the DB like so:

async with get_session() as session:
    role = UserRole(description=UserRoleType.client)
    session.add(role)
    await session.commit()

    user = User(username='test', user_role_id=role.id)
    session.add(user)
    await session.commit()
    await session.refresh(user)

And access the committed data with:

    results = await session.execute(select(User).where(User.id == 1)).one()

Output: (User(user_role_id=1, username='test', id=1),)

Notice that there's an user_role_id, but where's the user_role object?

In fact, if I try to access it, it raises an error:

AttributeError: Could not locate column in row for column 'user_role'

I also tried to pass the role instead of the user_role_id at the insertion of the User:

    ...
    user = User(username='test', user_role=role)

But I got the following error:

sqlalchemy.exc.InterfaceError: (sqlite3.InterfaceError) Error binding parameter 2 - probably unsupported type.

1

There are 1 best solutions below

0
On

A few things first.

  1. You did not include your import statements, so I will have to guess a few things.
  2. You probably want the User.user_role_id and User.user_role fields to be "pydantically" optional. This allows you to create user instances without passing the role to the constructor, giving you the option to do so after initialization or for example by appending User objects to the UserRole.users list instead. To enforce that a user must have a role on the database level, you simply define nullable=False on the User.user_role_id field. That way, if you try to commit to the DB without having defined a user role for a user in any of the possible ways, you will get an error.
  3. In your database insertion code you write role = UserRole(description=UserRoleType.client). I assume the description is from older code and you meant to write role = UserRole(type=UserRoleType.client).
  4. You probably want your UserRole.type to be not nullable on the database side. You can do so by passing nullable=False to the Column constructor (not the Field constructor).
  5. I will simplify a little by using blocking code (non-async) and a SQLite database.

This is how I would suggest you define the models:

from enum import Enum as EnumPy
from typing import Optional
from sqlalchemy.sql.schema import Column
from sqlalchemy.sql.sqltypes import Enum as EnumSQL
from sqlmodel import Field, Relationship, SQLModel


class UserRoleType(str, EnumPy):
    admin = 'admin'
    client = 'client'


class UserRole(SQLModel, table=True):
    __tablename__ = 'user_role'

    id: Optional[int] = Field(default=None, primary_key=True)
    type: UserRoleType = Field(
        default=UserRoleType.client,
        sa_column=Column(EnumSQL(UserRoleType), nullable=False),
    )
    write_access: bool = Field(default=False)
    read_access: bool = Field(default=False)

    users: list["User"] = Relationship(back_populates='user_role')


class User(SQLModel, table=True):
    __tablename__ = 'user'

    id: Optional[int] = Field(default=None, primary_key=True)
    username: str = Field(..., index=True)

    user_role_id: Optional[int] = Field(
        foreign_key='user_role.id',
        default=None,
        nullable=False,
    )
    user_role: Optional[UserRole] = Relationship(back_populates='users')

Here is a little test function to show that it works as expected:

from sqlmodel import Session, SQLModel, create_engine

# ... import User, UserRole, UserRoleType


def test() -> None:
    # Initialize database & session:
    engine = create_engine('sqlite:///', echo=True)
    SQLModel.metadata.drop_all(engine)
    SQLModel.metadata.create_all(engine)
    session = Session(engine)

    # Create the test objects:
    role = UserRole(type=UserRoleType.client)
    user = User(username='test', user_role=role)
    session.add(user)  # notice we don't need to add the `role` explicitly
    session.commit()
    session.refresh(user)

    # Do some checks:
    assert user.user_role.type is UserRoleType.client
    assert user.user_role_id == role.id and isinstance(role.id, int)
    assert role.users == [user]


if __name__ == '__main__':
    test()

This passes without errors. The SQL output generated by the engine is as follows:

CREATE TABLE user_role (
    type VARCHAR(6) NOT NULL, 
    id INTEGER NOT NULL, 
    write_access BOOLEAN NOT NULL, 
    read_access BOOLEAN NOT NULL, 
    PRIMARY KEY (id)
)

CREATE TABLE user (
    id INTEGER NOT NULL, 
    username VARCHAR NOT NULL, 
    user_role_id INTEGER NOT NULL, 
    PRIMARY KEY (id), 
    FOREIGN KEY(user_role_id) REFERENCES user_role (id)
)

CREATE INDEX ix_user_username ON user (username)

INSERT INTO user_role (type, write_access, read_access) VALUES ('client', 0, 0)

INSERT INTO user (username, user_role_id) VALUES ('test', 1)

SELECT user.id, user.username, user.user_role_id 
FROM user 
WHERE user.id = 1

SELECT user_role.type AS user_role_type, user_role.id AS user_role_id, user_role.write_access AS user_role_write_access, user_role.read_access AS user_role_read_access 
FROM user_role 
WHERE user_role.id = 1

SELECT user.id AS user_id, user.username AS user_username, user.user_role_id AS user_user_role_id 
FROM user 
WHERE 1 = user.user_role_id

PS: I know the question was posted a while ago, but maybe this still helps or helps someone else.