feat: Add support for SQLModel - DIA-65344 (#114)

## Problem

* FastAPI-SQLA is not compatible with [SQLModel](https://sqlmodel.tiangolo.com/).

## Solution

* Instantiate an SQLModel session if sqlmodel can be imported.
* SQLAlchemy's `sessionmaker` expect a `class_` argument and because
  SQLModel's Session inherits from SQLAlchemy's Session, it just works.

Also:
* Move `fastapi_sqla.models.Base` to `fastapi_sqla.sqla`: It is not a
  model but the SQLAlchemy declarative base.

Author: hadrien

.circleci/config.yml

@@ -55,10 +55,10 @@ workflows:
           matrix:
             parameters:
               python_version: ["3.9", "3.10", "3.11"]
-              sqlalchemy_version: ["1.4", "2.0"]
+              sqlalchemy_version: ["1.4", "2.0", "2.0-sqlmodel"]
               asyncpg: ["asyncpg", "noasyncpg"]
               aws_rds_iam: ["aws_rds_iam", "noaws_rds_iam"]
-              pydantic_version: ["1.10", "2.0", "2.1"]
+              pydantic_version: ["1", "2"]
 
       - release/release:
           name: release
@@ -97,7 +97,7 @@ jobs:
           Specify which version of python to run the tests against
       sqlalchemy_version:
         type: enum
-        enum: ["2.0", "1.4", "1.3"]
+        enum: ["2.0-sqlmodel", "2.0", "1.4", "1.3"]
         description: |
           Specify which version of sqlalchemy to run the tests against
       asyncpg:
@@ -110,7 +110,7 @@ jobs:
         description: To run tests with and without asyncpg installed.
       pydantic_version:
         type: enum
-        enum: ["1.10", "2.0", "2.1"]
+        enum: ["1", "2"]
     executor:
       name: python-postgres
       python_version: <<parameters.python_version>>

README.md

@@ -8,7 +8,7 @@
 
 
 Fastapi-SQLA is an [SQLAlchemy] extension for [FastAPI] easy to setup with support for
-pagination, asyncio, and [pytest].
+pagination, asyncio, [SQLModel] and [pytest].
 It supports SQLAlchemy>=1.3 and is fully compliant with [SQLAlchemy 2.0].
 It is developped, maintained and used on production by the team at [@dialoguemd] with
 love from Montreal 🇨🇦.
@@ -539,6 +539,41 @@ async def async_all_users_alt(
     return await paginate(select(User))
 ```
 
+# SQLModel support 🎉
+
+If your project uses [SQLModel], then `Session` dependency is an SQLModel session::
+
+```python
+    from http import HTTPStatus
+
+    from fastapi import FastAPI, HTTPException
+    from fastapi_sqla import Item, Page, Paginate, Session, setup
+    from sqlmodel import Field, SQLModel, select
+
+    class Hero(SQLModel, table=True):
+        id: int | None = Field(default=None, primary_key=True)
+        name: str
+        secret_name: str
+        age: int | None = None
+
+
+    app = FastAPI()
+    setup(app)
+
+    @app.get("/heros", response_model=Page[Hero])
+    def list_hero(paginate: Paginate) -> Page[Hero]:
+        return paginate(select(Hero))
+
+
+    @app.get("/heros/{hero_id}", response_model=Item[Hero])
+    def get_hero(hero_id: int, session: Session) -> Item[Hero]:
+        hero = session.get(Hero, hero_id)
+        if hero is None:
+            raise HTTPException(HTTPStatus.NOT_FOUND)
+        return {"data": hero}
+
+```
+
 # Pytest fixtures
 
 This library provides a set of utility fixtures, through its PyTest plugin, which is
@@ -709,6 +744,7 @@ $ poetry run tox
 [FastAPI background tasks]: https://fastapi.tiangolo.com/tutorial/background-tasks/
 [SQLAlchemy]: http://sqlalchemy.org/
 [SQLAlchemy 2.0]: https://docs.sqlalchemy.org/en/20/changelog/migration_20.html
+[SQLModel]: https://sqlmodel.tiangolo.com
 [`asyncpg`]: https://magicstack.github.io/asyncpg/current/
 [scalars]: https://docs.sqlalchemy.org/en/20/core/connections.html#sqlalchemy.engine.Result.scalars
 [alembic]: https://alembic.sqlalchemy.org/

fastapi_sqla/__init__.py

@@ -1,7 +1,13 @@
 from fastapi_sqla.base import setup
-from fastapi_sqla.models import Base, Collection, Item, Page
+from fastapi_sqla.models import Collection, Item, Page
 from fastapi_sqla.pagination import Paginate, PaginateSignature, Pagination
-from fastapi_sqla.sqla import Session, SessionDependency, SqlaSession, open_session
+from fastapi_sqla.sqla import (
+    Base,
+    Session,
+    SessionDependency,
+    SqlaSession,
+    open_session,
+)
 
 __all__ = [
     "Base",

fastapi_sqla/async_pagination.py

@@ -19,7 +19,7 @@
 
 
 async def default_query_count(session: SqlaAsyncSession, query: Select) -> int:
-    result = await session.execute(select(func.count()).select_from(query.subquery()))
+    result = await session.execute(select(func.count()).select_from(query.subquery()))  # type: ignore # noqa
     return cast(int, result.scalar())
 
 
@@ -34,10 +34,10 @@ async def paginate_query(
 ) -> Page:
     total_pages = math.ceil(total_items / limit)
     page_number = offset / limit + 1
-    query = query.offset(offset).limit(limit)
+    query = query.offset(offset).limit(limit)  # type: ignore
     result = await session.execute(query)
     data = iter(
-        cast(Iterator, result.unique().scalars() if scalars else result.mappings())
+        cast(Iterator, result.unique().scalars() if scalars else result.mappings())  # type: ignore # noqa
     )
     return Page(
         data=data,

fastapi_sqla/async_sqla.py

@@ -11,8 +11,7 @@
 from sqlalchemy.orm.session import sessionmaker
 
 from fastapi_sqla import aws_aurora_support, aws_rds_iam_support
-from fastapi_sqla.models import Base
-from fastapi_sqla.sqla import _DEFAULT_SESSION_KEY, new_engine
+from fastapi_sqla.sqla import _DEFAULT_SESSION_KEY, Base, new_engine
 
 logger = structlog.get_logger(__name__)
 

fastapi_sqla/models.py

@@ -2,14 +2,6 @@
 
 from pydantic import BaseModel, Field
 from pydantic import __version__ as pydantic_version
-from sqlalchemy.ext.declarative import DeferredReflection
-
-try:
-    from sqlalchemy.orm import DeclarativeBase
-except ImportError:
-    from sqlalchemy.ext.declarative import declarative_base
-
-    DeclarativeBase = declarative_base()  # type: ignore
 
 major, _, _ = [int(v) for v in pydantic_version.split(".")]
 is_pydantic2 = major == 2
@@ -19,10 +11,6 @@
     from pydantic.generics import GenericModel  # type:ignore
 
 
-class Base(DeclarativeBase, DeferredReflection):
-    __abstract__ = True
-
-
 ItemT = TypeVar("ItemT")
 
 

fastapi_sqla/pagination.py

@@ -68,7 +68,7 @@ def _paginate_legacy(
     total_pages = math.ceil(total_items / limit)
     page_number = offset / limit + 1
     return Page(
-        data=query.offset(offset).limit(limit).all(),
+        data=query.offset(offset).limit(limit).all(),  # type: ignore
         meta={
             "offset": offset,
             "total_items": total_items,
@@ -90,10 +90,10 @@ def _paginate(
 ) -> Page:
     total_pages = math.ceil(total_items / limit)
     page_number = offset / limit + 1
-    query = query.offset(offset).limit(limit)
+    query = query.offset(offset).limit(limit)  # type: ignore
     result = session.execute(query)
     data = iter(
-        cast(Iterator, result.unique().scalars() if scalars else result.mappings())
+        cast(Iterator, result.unique().scalars() if scalars else result.mappings())  # type: ignore  # noqa
     )
     return Page(
         data=data,

fastapi_sqla/sqla.py

@@ -2,19 +2,33 @@
 import os
 from collections.abc import Generator
 from contextlib import contextmanager
-from typing import Annotated
+from typing import Annotated, Generic, TypeVar
 
 import structlog
 from fastapi import Depends, Request
 from fastapi.concurrency import contextmanager_in_threadpool
 from fastapi.responses import PlainTextResponse
 from sqlalchemy import engine_from_config, text
 from sqlalchemy.engine import Engine
+from sqlalchemy.ext.declarative import DeferredReflection
 from sqlalchemy.orm.session import Session as SqlaSession
 from sqlalchemy.orm.session import sessionmaker
 
 from fastapi_sqla import aws_aurora_support, aws_rds_iam_support
-from fastapi_sqla.models import Base
+
+try:
+    from sqlalchemy.orm import DeclarativeBase
+except ImportError:
+    from sqlalchemy.ext.declarative import declarative_base
+
+    DeclarativeBase = declarative_base()  # type: ignore
+
+try:
+    from sqlmodel import Session as SqlaSession  # type: ignore  # noqa
+
+except ImportError:
+    pass
+
 
 logger = structlog.get_logger(__name__)
 
@@ -23,6 +37,10 @@
 _session_factories: dict[str, sessionmaker] = {}
 
 
+class Base(DeclarativeBase, DeferredReflection):
+    __abstract__ = True
+
+
 def new_engine(key: str = _DEFAULT_SESSION_KEY) -> Engine:
     envvar_prefix = "sqlalchemy_"
     if key != _DEFAULT_SESSION_KEY:
@@ -50,7 +68,8 @@ def startup(key: str = _DEFAULT_SESSION_KEY):
         raise
 
     Base.prepare(engine)
-    _session_factories[key] = sessionmaker(bind=engine)
+
+    _session_factories[key] = sessionmaker(bind=engine, class_=SqlaSession)
 
     logger.info("engine startup", engine_key=key, engine=engine)
 
@@ -146,11 +165,14 @@ def get_users(session: fastapi_sqla.Session):
     return response
 
 
-class SessionDependency:
+S = TypeVar("S", bound=SqlaSession)
+
+
+class SessionDependency(Generic[S]):
     def __init__(self, key: str = _DEFAULT_SESSION_KEY) -> None:
         self.key = key
 
-    def __call__(self, request: Request) -> SqlaSession:
+    def __call__(self, request: Request) -> S:
         """Yield the sqlalchemy session for that request.
 
         It is meant to be used as a FastAPI dependency::
@@ -175,5 +197,5 @@ def get_users(session: SqlaSession = Depends(SessionDependency())):
             raise
 
 
-default_session_dep = SessionDependency()
+default_session_dep = SessionDependency[SqlaSession]()
 Session = Annotated[SqlaSession, Depends(default_session_dep)]