SQLAlchemy typing support on filter for imperatively mapped attrs classes

457 views Asked by At

For my application I am trying to separate my "domain layer" classes from the database backend, in order to be able to unit test these classes independently of the database. I use mypy for static analysis and I get typing errors when using Query.filter() with the attributes of my imperatively mapped class. I don't want to use # type: ignore everywhere I use Query.filter() so I am looking for a correct solution to this problem.

I created the following model:

# project/models/user.py


from attrs import define, field


@define(slots=False, kw_only=True)
class User:
    id: int = field(init=False)
    username: str = field()
    hashed_password: str = field()
    is_active: bool = field(default=True)
    is_superuser: bool = field(default=False)

In production I of course use the database with the following table:

# project/db/tables/user.py

import sqlalchemy as sa
from sqlalchemy.sql import expression

from ..base import mapper_registry


user = sa.Table(
    "users",
    mapper_registry.metadata,
    sa.Column("id", sa.Integer, primary_key=True),
    sa.Column("username", sa.String(155), unique=True, nullable=False),
    sa.Column("hashed_password", sa.String(200), nullable=False),
    sa.Column("is_active", sa.Boolean, nullable=False, server_default=expression.true()),
    sa.Column(
        "is_superuser", sa.Boolean, nullable=False, server_default=expression.false()
    ),
)

And to map my class to the table I did the following:

# project/db/__init__.py

from .base import mapper_registry
from .tables import user

from project.models import User


mapper_registry.map_imperatively(User, user)

The functionality works as expected, correct queries are being constructed, however when I try to call Query.filter() with the attributes of my model class, I get typing errors from mypy, because Query.filter() expects column expressions and not plain python types (I use mypy bundled with sqlalchemy sqlalchemy = {version = "^2.0.15", extras = ["mypy"]} and I included the mypy plugin in mypy.ini)

# project/queries.py

from typing import List

from sqlalchemy import true
from sqlalchemy.orm import Session

from project.models import User


def list_active(session: Session) -> List[User]:
    return session.query(User).filter(User.is_active == true()).all()

And the error:

project/queries.py12: error: Argument 1 to "filter" of "Query" has incompatible type "bool"; expected "Union[ColumnElement[bool], _HasClauseElement, SQLCoreOperations[bool], ExpressionElementRole[bool], Callable[[], ColumnElement[bool]], LambdaElement]"  [arg-type]
1

There are 1 answers

0
Victor Sandoval On

TL;DR: No "correct" solution, in my opinion, because it is not referred in the documentation.

session.query(User).filter(users_table.columns["is_active"].is_(True)).all()

Also, I do not think the imperative mapping plays a role either, the story would be the same with declarative mapping (I use this style in my projects and have seen the same problems).

Explanation:

SQLAlchemy does not cater for this type check of mypy and there is no clean workaround as far as I know. Here is a piece of code to highlight some thoughts and an "alternative", though:

from typing import List
import sqlalchemy as sa
from sqlalchemy.orm import Session, registry
from sqlalchemy.sql.expression import BinaryExpression
from attrs import define, field

from sqlalchemy.sql import expression

engine = sa.create_engine("sqlite:///test.sqlite3")
mapper_registry = registry()


# Data model and table objects
@define(slots=False, kw_only=True)
class User:
    id: int = field(init=False)
    username: str = field()
    hashed_password: str = field()
    is_active: bool = field(default=True)
    is_superuser: bool = field(default=False)


user = sa.Table(
    "users",
    mapper_registry.metadata,
    sa.Column("id", sa.Integer, primary_key=True),
    sa.Column("username", sa.String(155), unique=True, nullable=False),
    sa.Column("hashed_password", sa.String(200), nullable=False),
    sa.Column(
        "is_active", sa.Boolean, nullable=False, server_default=expression.true()
    ),
    sa.Column(
        "is_superuser", sa.Boolean, nullable=False, server_default=expression.false()
    ),
)


# Map the User class to the 'users' table imperatively
mapper_registry.map_imperatively(User, user)


# Function to populate the database with initial data
def populate_database():
    # 'users' table creation
    user.create(engine)

    with Session(engine) as session:
        user_1 = User(username="John Doe", hashed_password="aaa")
        user_2 = User(username="Jane Smith", hashed_password="bbb")
        user_3 = User(username="Bob Johnson", hashed_password="ccc")

        session.add(user_1)
        session.add(user_2)
        session.add(user_3)
        session.commit()


def list_active() -> List[User]:
    with Session(engine) as session:
        expr: bool = User.is_active == sa.true()
        result = session.query(User).filter(expr).all()
    return result


def list_active_with_idiom() -> List[User]:
    with Session(engine) as session:
        expr: BinaryExpression = User.is_active.is_(True)
        result = session.query(User).filter(expr).all()
    return result


def list_active_with_types() -> List[User]:
    with Session(engine) as session:
        expr: BinaryExpression = user.columns["is_active"].is_(True)
        result = session.query(User).filter(user.columns["is_active"].is_(True)).all()
    return result


# Run this once only, to populate the table with initial data using the models
# populate_database()

# Run after populating the database
print(list_active())  # question way
print(list_active_with_idiom())  # SQLAlchemy recommended way
print(list_active_with_types())  # type checking sound

The snippet uses the original data models and Table object of your example, just put together in the same file to allow reproducing the results. After initialising the SQLite DB by calling populate_database() you can run the 3 alternative functions to list the active users. list_active() is your original function, list_active_with_idiom() is the one that uses SQLAlchemy's recommended way and the last one is the type-checking sound one.

  • list_active() evaluates the Python expression and creates a bool object that is fed into the filter() function, which is only annotated for _ColumnExpressionArgument[bool] (the type that shows in the mypy error message). This is the original one.

  • list_active_with_idiom() uses the SQLalchemy way, shown in the docs. It creates what is called a BinaryExpression (you can see it imported for typing and explanation purposes) out of a Column object, the True keyword and the is_() operator. These are, respectively, the left, right and operator attributes of the BinaryExpression. The base class of this class is OperatorExpression, which has ColumnElement as base class (making the type checking pass). You can notice, though, that the line will still generate a type checking problem as it relies on the is_() operator creating this BinaryExpression out of the data model attribute (which is a boolean). User.is_active.is_(True) works because of the mapping between the data models and the tables, but User.is_active is still just a boolean to mypy.

  • list_active_with_types() is the "type checking sound" version, and leverages what you can learn from the behaviour in the second option. users_table.columns["is_active"].is_(True) first grabs the Column object from the original Table object directly and then applies the same operator to compare to True. I am not sure this is how it is supposed to work, but makes mypy happy.

If you run this example, you should see the 2 following mypy errors:

main.py:61: error: Argument 1 to "filter" of "Query" has incompatible type "bool"; expected "Union[ColumnElement[bool], _HasClauseElement, SQLCoreOperations[bool], ExpressionElementRole[bool], Callable[[], ColumnElement[bool]], LambdaElement]"  [arg-type]
main.py:67: error: "bool" has no attribute "is_"  [attr-defined]
Found 2 errors in 1 file (checked 1 source file)

As you can see, no type checking error for the last option. Not sure this is quite right, but I have not found it anywhere in official documentation and I see it more as a workaround to SQLAlchemy's annotations.

Bonus comment: Projects like SQLModel use SQLAlchemy for the ORM part; among other things this project adds a good bit of annotations. Having a read of the source it looks like they skip the mypy type checking for some tricky lines like this one.