fix: allow paginating queries that return non scalars (#36)

Author: hadrien

README.md

@@ -143,6 +143,9 @@ class UserModel(BaseModel):
     id: int
     name: str
 
+    class Config:
+        orm_mode = True
+
 
 @router.get("/users", response_model=Page[UserModel])
 def all_users(paginate: Paginate = Depends()):
@@ -177,6 +180,47 @@ By default:
     }
     ```
 
+#### Paginating non-scalar results
+
+To paginate a query which doesn't return [scalars], specify `scalars=False` when invoking
+`paginate`:
+
+```python
+from fastapi import APIRouter, Depends
+from fastapi_sqla import Base, Page, Paginate
+from pydantic import BaseModel
+from sqlalchemy import func, select
+from sqlalchemy.orm import relationship
+
+router = APIRouter()
+
+
+class User(Base):
+    __tablename__ = "user"
+    notes = relationship("Note")
+
+
+class Note(Base):
+    __tablename__ = "note"
+
+
+class UserModel(BaseModel):
+    id: int
+    name: str
+    notes_count: int
+
+
+@router.get("/users", response_model=Page[UserModel])
+def all_users(paginate: Paginate = Depends()):
+    query = (
+        select(User.id, User.name, func.count(Note.id).label("notes_count"))
+        .join(Note)
+        .group_by(User)
+    )
+    return paginate(query, scalars=False)
+```
+
+
 #### Customize pagination
 
 You can customize:
@@ -359,3 +403,4 @@ It returns the path of  `alembic.ini` configuration file. By default, it returns
 [FastAPI background tasks]: https://fastapi.tiangolo.com/tutorial/background-tasks/
 [SQLAlchemy]: http://sqlalchemy.org/
 [`asyncpg`]: https://magicstack.github.io/asyncpg/current/
+[scalars]: (https://docs.sqlalchemy.org/en/14/core/connections.html?highlight=scalars#sqlalchemy.engine.Result.scalars),

fastapi_sqla/__init__.py

@@ -210,6 +210,7 @@ def paginate_query(
     total_items: int,
     offset: int,
     limit: int,
+    scalars: bool = True,
 ) -> Page[T]:  # pragma no cover
     "Dispatch on registered functions based on `query` type"
     raise NotImplementedError(f"no paginate_query registered for type {type(query)!r}")
@@ -222,6 +223,7 @@ def _paginate_legacy(
     total_items: int,
     offset: int,
     limit: int,
+    scalars: bool = True,
 ) -> Page[T]:
     total_pages = math.ceil(total_items / limit)
     page_number = offset / limit + 1
@@ -243,13 +245,16 @@ def _paginate(
     total_items: int,
     offset: int,
     limit: int,
+    *,
+    scalars: bool = True,
 ) -> Page[T]:
     total_pages = math.ceil(total_items / limit)
     page_number = offset / limit + 1
     query = query.offset(offset).limit(limit)
     result = session.execute(query)
+    data = iter(result.unique().scalars() if scalars else result.unique().mappings())
     return Page[T](
-        data=iter(result.unique().scalars()),
+        data=data,
         meta={
             "offset": offset,
             "total_items": total_items,
@@ -269,9 +274,11 @@ def dependency(
         offset: int = Query(0, ge=0),
         limit: int = Query(min_page_size, ge=1, le=max_page_size),
     ) -> PaginateSignature:
-        def paginate(query: DbQuery) -> Page[T]:
+        def paginate(query: DbQuery, scalars=True) -> Page[T]:
             total_items = query_count(session, query)
-            return paginate_query(query, session, total_items, offset, limit)
+            return paginate_query(
+                query, session, total_items, offset, limit, scalars=scalars
+            )
 
         return paginate
 

tests/test_pagination.py

@@ -188,6 +188,11 @@ class User(BaseModel):
         class Config:
             orm_mode = True
 
+    class UserWithNotesCount(BaseModel):
+        id: int
+        name: str
+        notes_count: int
+
     @app.get("/v1/users", response_model=Page[User])
     def sqla_13_all_users(session: Session = Depends(), paginate: Paginate = Depends()):
         query = (
@@ -200,6 +205,18 @@ def sqla_14_all_users(paginate: Paginate = Depends()):
         query = select(user_cls).options(joinedload("notes")).order_by(user_cls.id)
         return paginate(query)
 
+    @app.get("/v2/users-with-notes-count", response_model=Page[UserWithNotesCount])
+    def sqla_14_all_users_with_notes_count(paginate: Paginate = Depends()):
+        query = (
+            select(
+                user_cls.id, user_cls.name, func.count(note_cls.id).label("notes_count")
+            )
+            .join(note_cls)
+            .order_by(user_cls.id)
+            .group_by(user_cls)
+        )
+        return paginate(query, scalars=False)
+
     def query_count(session: Session, query: Select) -> int:
         return session.execute(select(func.count()).select_from(user_cls)).scalar()
 
@@ -236,6 +253,9 @@ async def client(app):
         param(0, 10, "/v2/custom/users", marks=mark.sqlalchemy("1.4")),
         param(10, 10, "/v2/custom/users", marks=mark.sqlalchemy("1.4")),
         param(40, 2, "/v2/custom/users", marks=mark.sqlalchemy("1.4")),
+        param(0, 10, "/v2/users-with-notes-count", marks=mark.sqlalchemy("1.4")),
+        param(10, 10, "/v2/users-with-notes-count", marks=mark.sqlalchemy("1.4")),
+        param(40, 2, "/v2/users-with-notes-count", marks=mark.sqlalchemy("1.4")),
     ],
 )
 async def test_functional(client, offset, items_number, path):