Merge pull request #22 from febus982/repository_pagination

Pagination and protocols

Author: febus982

.coveragerc

@@ -8,3 +8,4 @@ parallel = true
 exclude_lines =
     pragma: no cover
     pass
+    \.\.\.

README.md

@@ -198,6 +198,7 @@ The classes provide some common use methods:
 * `save_many`: Persist multiple models in a single transaction
 * `delete`: Delete a model
 * `find`: Search for a list of models (basically an adapter for SELECT queries)
+* `paginated_find`: Search for a list of models, with pagination support
 
 ### Session lifecycle in repositories
 

sqlalchemy_bind_manager/__init__.py

@@ -7,6 +7,7 @@
     SQLAlchemyRepository,
     SQLAlchemyAsyncRepository,
     SortDirection,
+    PaginatedResult,
 )
 from ._unit_of_work import (
     UnitOfWork,

sqlalchemy_bind_manager/_repository/__init__.py

@@ -1,3 +1,3 @@
 from .sync import SQLAlchemyRepository
 from .async_ import SQLAlchemyAsyncRepository
-from .common import SortDirection
+from .common import SortDirection, PaginatedResult

sqlalchemy_bind_manager/_repository/async_.py

@@ -1,14 +1,23 @@
 from abc import ABC
 from contextlib import asynccontextmanager
-from typing import Union, Generic, Tuple, Iterable, List, AsyncIterator, Any, Mapping
+from typing import (
+    Union,
+    Generic,
+    Tuple,
+    Iterable,
+    List,
+    AsyncIterator,
+    Any,
+    Mapping,
+    Type,
+)
 
-from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from .._bind_manager import SQLAlchemyAsyncBind
 from .._transaction_handler import AsyncSessionHandler
 from ..exceptions import ModelNotFound, InvalidConfig
-from .common import MODEL, PRIMARY_KEY, SortDirection, BaseRepository
+from .common import MODEL, PRIMARY_KEY, SortDirection, BaseRepository, PaginatedResult
 
 
 class SQLAlchemyAsyncRepository(Generic[MODEL], BaseRepository[MODEL], ABC):
@@ -19,14 +28,17 @@ def __init__(
         self,
         bind: Union[SQLAlchemyAsyncBind, None] = None,
         session: Union[AsyncSession, None] = None,
+        model_class: Union[Type[MODEL], None] = None,
     ) -> None:
         """
         :param bind: A configured instance of SQLAlchemyAsyncBind
-        :type bind: SQLAlchemyAsyncBind
+        :type bind: Union[SQLAlchemyAsyncBind, None]
         :param session: An externally managed session
-        :type session: AsyncSession
+        :type session: Union[AsyncSession, None]
+        :param model_class: A mapped SQLAlchemy model
+        :type model_class: Union[Type[MODEL], None]
         """
-        super().__init__()
+        super().__init__(model_class=model_class)
         if not (bool(bind) ^ bool(session)):
             raise InvalidConfig("Either `bind` or `session` have to be used, not both")
         self._external_session = session
@@ -105,16 +117,55 @@ async def find(
 
         :param order_by:
         :param search_params: A dictionary containing equality filters
+        :param limit: Number of models to retrieve
+        :type limit: int
+        :param offset: Number of models to skip
+        :type offset: int
         :return: A collection of models
         :rtype: List
         """
-        stmt = select(self._model)
-        if search_params:
-            stmt = self._filter_select(stmt, search_params)
-
-        if order_by is not None:
-            stmt = self._filter_order_by(stmt, order_by)
+        stmt = self._find_query(search_params, order_by)
 
         async with self._get_session() as session:
             result = await session.execute(stmt)
             return [x for x in result.scalars()]
+
+    async def paginated_find(
+        self,
+        per_page: int,
+        page: int,
+        search_params: Union[None, Mapping[str, Any]] = None,
+        order_by: Union[None, Iterable[Union[str, Tuple[str, SortDirection]]]] = None,
+    ) -> PaginatedResult[MODEL]:
+        """Find models using filters and pagination
+
+        E.g.
+        find(name="John") finds all models with name = John
+
+        :param per_page: Number of models to retrieve
+        :type per_page: int
+        :param page: Page to retrieve
+        :type page: int
+        :param search_params: A dictionary containing equality filters
+        :param order_by:
+        :return: A collection of models
+        :rtype: List
+        """
+
+        find_stmt = self._find_query(search_params, order_by)
+        paginated_stmt = self._paginate_query(find_stmt, page, per_page)
+
+        async with self._get_session() as session:
+            total_items_count = (
+                await session.execute(self._count_query(find_stmt))
+            ).scalar() or 0
+            result_items = [
+                x for x in (await session.execute(paginated_stmt)).scalars()
+            ]
+
+            return self._build_paginated_result(
+                result_items=result_items,
+                total_items_count=total_items_count,
+                page=page,
+                per_page=per_page,
+            )

sqlalchemy_bind_manager/_repository/common.py

@@ -1,10 +1,23 @@
 from abc import ABC
 from enum import Enum
 from functools import partial
-from typing import TypeVar, Union, Generic, Type, Tuple, Iterable, Any, Mapping
-
-from sqlalchemy import asc, desc
-from sqlalchemy.orm import object_mapper, class_mapper, Mapper
+from math import ceil
+from typing import (
+    TypeVar,
+    Union,
+    Generic,
+    Type,
+    Tuple,
+    Iterable,
+    Any,
+    Mapping,
+    List,
+    Collection,
+)
+
+from pydantic.generics import GenericModel
+from sqlalchemy import asc, desc, select, func
+from sqlalchemy.orm import object_mapper, class_mapper, Mapper, lazyload
 from sqlalchemy.orm.exc import UnmappedInstanceError
 from sqlalchemy.sql import Select
 
@@ -19,9 +32,30 @@ class SortDirection(Enum):
     DESC = partial(desc)
 
 
+class PaginatedResult(GenericModel, Generic[MODEL]):
+    items: List[MODEL]
+    page: int
+    per_page: int
+    total_pages: int
+    total_items: int
+
+
 class BaseRepository(Generic[MODEL], ABC):
+    _max_query_limit: int = 50
     _model: Type[MODEL]
 
+    def __init__(self, model_class: Union[Type[MODEL], None] = None) -> None:
+        if getattr(self, "_model", None) is None and model_class is not None:
+            self._model = model_class
+
+        if getattr(self, "_model", None) is None or not self._is_mapped_object(
+            self._model()
+        ):
+            raise InvalidModel(
+                "You need to supply a valid model class either in the `model_class` parameter"
+                " or in the `_model` class property."
+            )
+
     def _is_mapped_object(self, obj: object) -> bool:
         """Checks if the object is handled by the repository and is mapped in SQLAlchemy.
 
@@ -65,7 +99,6 @@ def _filter_select(self, stmt: Select, search_params: Mapping[str, Any]) -> Sele
         :param search_params: Any keyword argument to be used as equality filter
         :return: The filtered query
         """
-        # TODO: Add support for offset/limit
         # TODO: Add support for relationship eager load
         for k, v in search_params.items():
             """
@@ -100,3 +133,75 @@ def _filter_order_by(
                 stmt = stmt.order_by(value[1].value(getattr(self._model, value[0])))
 
         return stmt
+
+    def _find_query(
+        self,
+        search_params: Union[None, Mapping[str, Any]] = None,
+        order_by: Union[None, Iterable[Union[str, Tuple[str, SortDirection]]]] = None,
+    ) -> Select:
+        stmt = select(self._model)
+
+        if search_params:
+            stmt = self._filter_select(stmt, search_params)
+        if order_by is not None:
+            stmt = self._filter_order_by(stmt, order_by)
+
+        return stmt
+
+    def _count_query(
+        self,
+        query: Select,
+    ) -> Select:
+        return select(func.count()).select_from(
+            query.options(lazyload("*")).order_by(None).subquery()  # type: ignore
+        )
+
+    def _paginate_query(
+        self,
+        stmt: Select,
+        page: int,
+        per_page: int,
+    ) -> Select:
+        """Build the query offset and limit clauses from submitted parameters.
+
+        :param stmt: a Select statement
+        :type stmt: Select
+        :param page: Number of models to skip
+        :type page: int
+        :param per_page: Number of models to retrieve
+        :type per_page: int
+        :return: The filtered query
+        """
+
+        _offset = max((page - 1) * per_page, 0)
+        if _offset > 0:
+            stmt = stmt.offset(_offset)
+
+        _limit = max(min(per_page, self._max_query_limit), 0)
+        stmt = stmt.limit(_limit)
+
+        return stmt
+
+    def _build_paginated_result(
+        self,
+        result_items: Collection[MODEL],
+        total_items_count: int,
+        page: int,
+        per_page: int,
+    ) -> PaginatedResult:
+
+        _per_page = max(min(per_page, self._max_query_limit), 0)
+        total_pages = (
+            0
+            if total_items_count == 0 or total_items_count is None
+            else ceil(total_items_count / _per_page)
+        )
+        _page = 0 if len(result_items) == 0 else min(page, total_pages)
+
+        return PaginatedResult(
+            items=result_items,
+            page=_page,
+            per_page=_per_page,
+            total_items=total_items_count,
+            total_pages=total_pages,
+        )

sqlalchemy_bind_manager/_repository/sync.py

@@ -1,14 +1,24 @@
 from abc import ABC
 from contextlib import contextmanager
-from typing import Union, Generic, Iterable, Tuple, List, Iterator, Any, Mapping
+from math import ceil
+from typing import (
+    Union,
+    Generic,
+    Iterable,
+    Tuple,
+    List,
+    Iterator,
+    Any,
+    Mapping,
+    Type,
+)
 
-from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from .._bind_manager import SQLAlchemyBind
 from .._transaction_handler import SessionHandler
 from ..exceptions import ModelNotFound, InvalidConfig
-from .common import MODEL, PRIMARY_KEY, SortDirection, BaseRepository
+from .common import MODEL, PRIMARY_KEY, SortDirection, BaseRepository, PaginatedResult
 
 
 class SQLAlchemyRepository(Generic[MODEL], BaseRepository[MODEL], ABC):
@@ -19,14 +29,17 @@ def __init__(
         self,
         bind: Union[SQLAlchemyBind, None] = None,
         session: Union[Session, None] = None,
+        model_class: Union[Type[MODEL], None] = None,
     ) -> None:
         """
         :param bind: A configured instance of SQLAlchemyBind
-        :type bind: SQLAlchemyBind
+        :type bind: Union[SQLAlchemyBind, None]
         :param session: An externally managed session
-        :type session: Session
+        :type session: Union[Session, None]
+        :param model_class: A mapped SQLAlchemy model
+        :type model_class: Union[Type[MODEL], None]
         """
-        super().__init__()
+        super().__init__(model_class=model_class)
         if not (bool(bind) ^ bool(session)):
             raise InvalidConfig("Either `bind` or `session` have to be used, not both")
         self._external_session = session
@@ -97,18 +110,51 @@ def find(
         E.g.
         find(name="John") finds all models with name = John
 
-        :param order_by:
         :param search_params: A dictionary containing equality filters
+        :param order_by:
         :return: A collection of models
         :rtype: List
         """
-        stmt = select(self._model)
-        if search_params:
-            stmt = self._filter_select(stmt, search_params)
-
-        if order_by is not None:
-            stmt = self._filter_order_by(stmt, order_by)
+        stmt = self._find_query(search_params, order_by)
 
         with self._get_session() as session:
             result = session.execute(stmt)
             return [x for x in result.scalars()]
+
+    def paginated_find(
+        self,
+        per_page: int,
+        page: int,
+        search_params: Union[None, Mapping[str, Any]] = None,
+        order_by: Union[None, Iterable[Union[str, Tuple[str, SortDirection]]]] = None,
+    ) -> PaginatedResult[MODEL]:
+        """Find models using filters and pagination
+
+        E.g.
+        find(name="John") finds all models with name = John
+
+        :param per_page: Number of models to retrieve
+        :type per_page: int
+        :param page: Page to retrieve
+        :type page: int
+        :param search_params: A dictionary containing equality filters
+        :param order_by:
+        :return: A collection of models
+        :rtype: List
+        """
+
+        find_stmt = self._find_query(search_params, order_by)
+        paginated_stmt = self._paginate_query(find_stmt, page, per_page)
+
+        with self._get_session() as session:
+            total_items_count = (
+                session.execute(self._count_query(find_stmt)).scalar() or 0
+            )
+            result_items = [x for x in session.execute(paginated_stmt).scalars()]
+
+            return self._build_paginated_result(
+                result_items=result_items,
+                total_items_count=total_items_count,
+                page=page,
+                per_page=per_page,
+            )