Add async pymongo integration (#1544)

Author: uriyyo

fastapi_pagination/ext/pymongo.py

@@ -1,17 +1,22 @@
 from __future__ import annotations
 
-__all__ = ["paginate"]
+__all__ = [
+    "apaginate",
+    "paginate",
+]
+
 
 from collections.abc import Mapping, Sequence
 from typing import Any, Optional, TypeVar
 
+from pymongo.asynchronous.collection import AsyncCollection
 from pymongo.collection import Collection
 
 from fastapi_pagination.bases import AbstractParams
 from fastapi_pagination.config import Config
-from fastapi_pagination.flow import flow_expr, run_sync_flow
+from fastapi_pagination.flow import flow_expr, run_async_flow, run_sync_flow
 from fastapi_pagination.flows import generic_flow
-from fastapi_pagination.types import AdditionalData, SyncItemsTransformer
+from fastapi_pagination.types import AdditionalData, ItemsTransformer, SyncItemsTransformer
 
 T = TypeVar("T", bound=Mapping[str, Any])
 
@@ -49,3 +54,38 @@ def paginate(
             config=config,
         )
     )
+
+
+async def apaginate(
+    collection: AsyncCollection[T],
+    query_filter: Optional[dict[Any, Any]] = None,
+    filter_fields: Optional[dict[Any, Any]] = None,
+    params: Optional[AbstractParams] = None,
+    sort: Optional[Sequence[Any]] = None,
+    *,
+    transformer: Optional[ItemsTransformer] = None,
+    additional_data: Optional[AdditionalData] = None,
+    config: Optional[Config] = None,
+    **kwargs: Any,
+) -> Any:
+    query_filter = query_filter or {}
+
+    return await run_async_flow(
+        generic_flow(
+            total_flow=flow_expr(lambda: collection.count_documents(query_filter)),
+            limit_offset_flow=flow_expr(
+                lambda raw_params: collection.find(
+                    query_filter,
+                    filter_fields,
+                    skip=raw_params.offset,
+                    limit=raw_params.limit,
+                    sort=sort,
+                    **kwargs,
+                ).to_list()
+            ),
+            params=params,
+            transformer=transformer,
+            additional_data=additional_data,
+            config=config,
+        )
+    )

tests/ext/test_pymongo.py

@@ -1,8 +1,10 @@
 import pytest
-from pymongo import MongoClient
+from pymongo import AsyncMongoClient, MongoClient
+from pytest_asyncio import fixture as async_fixture
 
-from fastapi_pagination.ext.pymongo import paginate
-from tests.base import BasePaginationTestSuite
+from fastapi_pagination.ext.pymongo import apaginate, paginate
+from tests.base import BasePaginationTestSuite, async_sync_testsuite
+from tests.utils import maybe_async
 
 from .utils import mongodb_test
 
@@ -12,18 +14,28 @@ def database_url(mongodb_url) -> str:
     return mongodb_url
 
 
-@pytest.fixture(scope="session")
-def db_client(database_url):
-    with MongoClient(database_url) as client:
-        yield client
-
-
 @mongodb_test
+@async_sync_testsuite
 class TestPymongo(BasePaginationTestSuite):
+    @async_fixture(scope="session")
+    async def db_client(self, database_url, is_async_db):
+        if is_async_db:
+            async with AsyncMongoClient(database_url) as client:
+                yield client
+        else:
+            with MongoClient(database_url) as client:
+                yield client
+
     @pytest.fixture(scope="session")
-    def app(self, builder, db_client):
+    def paginate_func(self, is_async_db):
+        return apaginate if is_async_db else paginate
+
+    @pytest.fixture(scope="session")
+    def app(self, builder, db_client, paginate_func):
+        builder = builder.new()
+
         @builder.both.default
         async def route():
-            return paginate(db_client.test.users)
+            return await maybe_async(paginate_func(db_client.test.users))
 
         return builder.build()