Add option to control sqlalchemy count query (#624)

Author: uriyyo

.gitignore

@@ -134,3 +134,6 @@ dmypy.json
 
 # IDEA
 .idea/
+
+# ruff
+.ruff_cache/

fastapi_pagination/ext/sqlalchemy.py

@@ -10,7 +10,7 @@
 from contextlib import suppress
 from typing import TYPE_CHECKING, Any, Optional, Tuple, Union, overload
 
-from sqlalchemy import func, literal_column, select
+from sqlalchemy import func, select
 from sqlalchemy.orm import Query, Session, noload
 from typing_extensions import TypeAlias
 
@@ -51,9 +51,16 @@ def paginate_query(query: Select, params: AbstractParams) -> Select:
     return generic_query_apply_params(query, params.to_raw_params().as_limit_offset())
 
 
-def count_query(query: Select) -> Select:
-    count_subquery = query.order_by(None).options(noload("*")).subquery()
-    return select(func.count(literal_column("*"))).select_from(count_subquery)
+def count_query(query: Select, *, use_subquery: bool = True) -> Select:
+    query = query.order_by(None).options(noload("*"))
+
+    if use_subquery:
+        return select(func.count()).select_from(query.subquery())
+
+    return query.with_only_columns(  # noqa: PIE804
+        func.count(),
+        **{"maintain_column_froms": True},
+    )
 
 
 def _maybe_unique(result: Any, unique: bool) -> Any:
@@ -66,6 +73,7 @@ def exec_pagination(
     conn: SyncConn,
     transformer: Optional[ItemsTransformer] = None,
     additional_data: AdditionalData = None,
+    subquery_count: bool = True,
     unique: bool = True,
     async_: bool = False,
 ) -> AbstractPage[Any]:
@@ -82,6 +90,8 @@ def _apply_items_transformer(*args: Any, **kwargs: Any) -> Any:
     if is_cursor(raw_params):
         if paging is None:
             raise ImportError("sqlakeyset is not installed")
+        if not getattr(query, "_order_by_clauses", True):
+            raise ValueError("Cursor pagination requires ordering")
 
         page = paging.select_page(
             conn,
@@ -100,7 +110,7 @@ def _apply_items_transformer(*args: Any, **kwargs: Any) -> Any:
             **(additional_data or {}),
         )
 
-    total = conn.scalar(count_query(query))
+    total = conn.scalar(count_query(query, use_subquery=subquery_count))
     query = paginate_query(query, params)
     items = _maybe_unique(conn.execute(query), unique)
     items = unwrap_scalars(items)
@@ -130,6 +140,7 @@ def paginate(
     query: Query[Any],
     params: Optional[AbstractParams] = None,
     *,
+    subquery_count: bool = True,
     transformer: Optional[SyncItemsTransformer] = None,
     additional_data: AdditionalData = None,
 ) -> Any:
@@ -142,6 +153,7 @@ def paginate(
     query: Select,
     params: Optional[AbstractParams] = None,
     *,
+    subquery_count: bool = True,
     transformer: Optional[SyncItemsTransformer] = None,
     additional_data: AdditionalData = None,
     unique: bool = True,
@@ -155,6 +167,7 @@ async def paginate(
     query: Select,
     params: Optional[AbstractParams] = None,
     *,
+    subquery_count: bool = True,
     transformer: Optional[AsyncItemsTransformer] = None,
     additional_data: AdditionalData = None,
     unique: bool = True,
@@ -166,9 +179,9 @@ def paginate(*args: Any, **kwargs: Any) -> Any:
     try:
         assert args
         assert isinstance(args[0], Query)
-        query, conn, params, transformer, additional_data, unique = _old_paginate_sign(*args, **kwargs)
+        query, conn, params, transformer, additional_data, unique, subquery_count = _old_paginate_sign(*args, **kwargs)
     except (TypeError, AssertionError):
-        query, conn, params, transformer, additional_data, unique = _new_paginate_sign(*args, **kwargs)
+        query, conn, params, transformer, additional_data, unique, subquery_count = _new_paginate_sign(*args, **kwargs)
 
     params, _ = verify_params(params, "limit-offset", "cursor")
 
@@ -181,20 +194,22 @@ def paginate(*args: Any, **kwargs: Any) -> Any:
             sync_conn,
             transformer,
             additional_data,
+            subquery_count,
             unique,
             async_=True,
         )
 
-    return exec_pagination(query, params, conn, transformer, additional_data, unique, async_=False)
+    return exec_pagination(query, params, conn, transformer, additional_data, subquery_count, unique, async_=False)
 
 
 def _old_paginate_sign(
     query: Query[Any],
     params: Optional[AbstractParams] = None,
     *,
+    subquery_count: bool = True,
     transformer: Optional[ItemsTransformer] = None,
     additional_data: AdditionalData = None,
-) -> Tuple[Select, SyncConn, Optional[AbstractParams], Optional[ItemsTransformer], AdditionalData, bool]:
+) -> Tuple[Select, SyncConn, Optional[AbstractParams], Optional[ItemsTransformer], AdditionalData, bool, bool]:
     if query.session is None:
         raise ValueError("query.session is None")
 
@@ -205,16 +220,17 @@ def _old_paginate_sign(
         stacklevel=3,
     )
 
-    return query, query.session, params, transformer, additional_data, True  # type: ignore
+    return query, query.session, params, transformer, additional_data, True, subquery_count  # type: ignore
 
 
 def _new_paginate_sign(
     conn: SyncConn,
     query: Select,
     params: Optional[AbstractParams] = None,
     *,
+    subquery_count: bool = True,
     transformer: Optional[ItemsTransformer] = None,
     additional_data: AdditionalData = None,
     unique: bool = True,
-) -> Tuple[Select, SyncConn, Optional[AbstractParams], Optional[ItemsTransformer], AdditionalData, bool]:
-    return query, conn, params, transformer, additional_data, unique
+) -> Tuple[Select, SyncConn, Optional[AbstractParams], Optional[ItemsTransformer], AdditionalData, bool, bool]:
+    return query, conn, params, transformer, additional_data, unique, subquery_count

fastapi_pagination/ext/sqlmodel.py

@@ -20,6 +20,7 @@ def paginate(
     query: Select[TSQLModel],
     params: Optional[AbstractParams] = None,
     *,
+    subquery_count: bool = True,
     transformer: Optional[SyncItemsTransformer] = None,
     additional_data: AdditionalData = None,
     unique: bool = True,
@@ -33,6 +34,7 @@ def paginate(
     query: SelectOfScalar[T],
     params: Optional[AbstractParams] = None,
     *,
+    subquery_count: bool = True,
     transformer: Optional[SyncItemsTransformer] = None,
     additional_data: AdditionalData = None,
     unique: bool = True,
@@ -46,6 +48,7 @@ def paginate(
     query: Type[TSQLModel],
     params: Optional[AbstractParams] = None,
     *,
+    subquery_count: bool = True,
     transformer: Optional[SyncItemsTransformer] = None,
     additional_data: AdditionalData = None,
     unique: bool = True,
@@ -59,6 +62,7 @@ async def paginate(
     query: Select[TSQLModel],
     params: Optional[AbstractParams] = None,
     *,
+    subquery_count: bool = True,
     transformer: Optional[AsyncItemsTransformer] = None,
     additional_data: AdditionalData = None,
     unique: bool = True,
@@ -72,6 +76,7 @@ async def paginate(
     query: SelectOfScalar[T],
     params: Optional[AbstractParams] = None,
     *,
+    subquery_count: bool = True,
     transformer: Optional[AsyncItemsTransformer] = None,
     additional_data: AdditionalData = None,
     unique: bool = True,
@@ -85,6 +90,7 @@ async def paginate(
     query: Type[TSQLModel],
     params: Optional[AbstractParams] = None,
     *,
+    subquery_count: bool = True,
     transformer: Optional[AsyncItemsTransformer] = None,
     additional_data: AdditionalData = None,
 ) -> Any:
@@ -97,6 +103,7 @@ def paginate(
     query: Any,
     params: Optional[AbstractParams] = None,
     *,
+    subquery_count: bool = True,
     transformer: Optional[ItemsTransformer] = None,
     additional_data: AdditionalData = None,
     unique: bool = True,
@@ -108,6 +115,7 @@ def paginate(
         session,
         query,
         params,
+        subquery_count=subquery_count,
         transformer=transformer,
         additional_data=additional_data,
         unique=unique,

tests/ext/test_sqlalchemy.py

@@ -1,18 +1,31 @@
 from typing import Iterator
 
 from fastapi import Depends, FastAPI
-from pytest import fixture
+from pytest import fixture, skip
+from sqlalchemy import select
 from sqlalchemy.orm.session import Session
 
 from fastapi_pagination import LimitOffsetPage, Page, add_pagination
 from fastapi_pagination.ext.sqlalchemy import paginate
 
 from ..base import BasePaginationTestCase
-from .utils import sqlalchemy20
+from .utils import is_sqlalchemy20, sqlalchemy20
+
+
+@fixture(
+    scope="session",
+    params=[True, False],
+    ids=["subquery_count", "no_subquery_count"],
+)
+def use_subquery_count(request):
+    if request.param and not is_sqlalchemy20:
+        skip("subquery_count is not supported for SQLAlchemy<2.0")
+
+    return request.param
 
 
 @fixture(scope="session")
-def app(sa_user, sa_session, model_cls):
+def app(sa_user, sa_session, model_cls, use_subquery_count):
     app = FastAPI()
 
     def get_db() -> Iterator[Session]:
@@ -25,7 +38,7 @@ def get_db() -> Iterator[Session]:
     @app.get("/default", response_model=Page[model_cls])
     @app.get("/limit-offset", response_model=LimitOffsetPage[model_cls])
     def route(db: Session = Depends(get_db)):
-        return paginate(db.query(sa_user))
+        return paginate(db, select(sa_user), subquery_count=use_subquery_count)
 
     return add_pagination(app)
 

tests/ext/test_sqlalchemy_cursor.py

@@ -1,5 +1,6 @@
 from typing import Iterator, List
 
+from _pytest.python_api import raises
 from fastapi import Depends, FastAPI, status
 from pydantic import parse_obj_as
 from pytest import fixture, mark
@@ -30,6 +31,10 @@ def get_db() -> Iterator[Session]:
     def route(db: Session = Depends(get_db)):
         return paginate(db, select(sa_user).order_by(sa_user.id, sa_user.name))
 
+    @app.get("/no-order", response_model=CursorPage[UserOut])
+    def route_on_order(db: Session = Depends(get_db)):
+        return paginate(db, select(sa_user))
+
     return add_pagination(app)
 
 
@@ -74,3 +79,13 @@ async def test_cursor(app, client, entities):
         cursor = data["previous_page"]
 
     assert items == entities
+
+
+@sqlalchemy20
+@mark.asyncio
+async def test_no_order(app, client, entities):
+    with raises(
+        ValueError,
+        match="^Cursor pagination requires ordering$",
+    ):
+        await client.get("/no-order")

tests/ext/test_sqlalchemy_deprecated.py

@@ -0,0 +1,35 @@
+from typing import Iterator
+
+from fastapi import Depends, FastAPI
+from pytest import fixture
+from sqlalchemy.orm.session import Session
+
+from fastapi_pagination import LimitOffsetPage, Page, add_pagination
+from fastapi_pagination.ext.sqlalchemy import paginate
+
+from ..base import BasePaginationTestCase
+from .utils import sqlalchemy20
+
+
+@fixture(scope="session")
+def app(sa_user, sa_session, model_cls):
+    app = FastAPI()
+
+    def get_db() -> Iterator[Session]:
+        db = sa_session()
+        try:
+            yield db
+        finally:
+            db.close()
+
+    @app.get("/default", response_model=Page[model_cls])
+    @app.get("/limit-offset", response_model=LimitOffsetPage[model_cls])
+    def route(db: Session = Depends(get_db)):
+        return paginate(db.query(sa_user))
+
+    return add_pagination(app)
+
+
+@sqlalchemy20
+class TestSQLAlchemy(BasePaginationTestCase):
+    pass

tests/ext/utils.py

@@ -1,3 +1,17 @@
 from pytest import mark
 
+try:
+    from sqlalchemy import __version__ as sqlalchemy_version
+
+    is_sqlalchemy20 = tuple(map(int, sqlalchemy_version.split("."))) >= (2, 0, 0)
+
+    del sqlalchemy_version
+except (ImportError, AttributeError):
+    is_sqlalchemy20 = False
+
 sqlalchemy20 = mark.sqlalchemy20
+
+only_sqlalchemy20 = mark.skipif(
+    lambda: not is_sqlalchemy20,
+    reason="Only for SQLAlchemy 2.0",
+)