add OffsetPaginationExtension extension (#757)

Author: vincentsarago

CHANGES.md

@@ -6,6 +6,10 @@
 
 * Removed `cql2-text` in supported `filter-lang` for `FilterExtensionPostRequest` model (as per specification)
 
+### Added
+
+* Add `OffsetPaginationExtension` extension to add `offset` query/body parameter to endpoints
+
 ## [3.0.2] - 2024-09-20
 
 ### Added

stac_fastapi/extensions/stac_fastapi/extensions/core/__init__.py

@@ -5,7 +5,11 @@
 from .fields import FieldsExtension
 from .filter import FilterExtension
 from .free_text import FreeTextAdvancedExtension, FreeTextExtension
-from .pagination import PaginationExtension, TokenPaginationExtension
+from .pagination import (
+    OffsetPaginationExtension,
+    PaginationExtension,
+    TokenPaginationExtension,
+)
 from .query import QueryExtension
 from .sort import SortExtension
 from .transaction import TransactionExtension
@@ -16,6 +20,7 @@
     "FilterExtension",
     "FreeTextExtension",
     "FreeTextAdvancedExtension",
+    "OffsetPaginationExtension",
     "PaginationExtension",
     "QueryExtension",
     "SortExtension",

stac_fastapi/extensions/stac_fastapi/extensions/core/pagination/__init__.py

@@ -1,6 +1,7 @@
 """Pagination classes as extensions."""
 
+from .offset_pagination import OffsetPaginationExtension
 from .pagination import PaginationExtension
 from .token_pagination import TokenPaginationExtension
 
-__all__ = ["PaginationExtension", "TokenPaginationExtension"]
+__all__ = ["OffsetPaginationExtension", "PaginationExtension", "TokenPaginationExtension"]

stac_fastapi/extensions/stac_fastapi/extensions/core/pagination/offset_pagination.py

@@ -0,0 +1,38 @@
+"""Pagination API extension."""
+
+from typing import List, Optional
+
+import attr
+from fastapi import FastAPI
+
+from stac_fastapi.types.extension import ApiExtension
+
+from .request import GETOffsetPagination, POSTOffsetPagination
+
+
+@attr.s
+class OffsetPaginationExtension(ApiExtension):
+    """Offset Pagination.
+
+    Though not strictly an extension, the chosen pagination will modify the form of the
+    request object. By making pagination an extension class, we can use
+    create_request_model to dynamically add the correct pagination parameter to the
+    request model for OpenAPI generation.
+    """
+
+    GET = GETOffsetPagination
+    POST = POSTOffsetPagination
+
+    conformance_classes: List[str] = attr.ib(factory=list)
+    schema_href: Optional[str] = attr.ib(default=None)
+
+    def register(self, app: FastAPI) -> None:
+        """Register the extension with a FastAPI application.
+
+        Args:
+            app: target FastAPI application.
+
+        Returns:
+            None
+        """
+        pass

stac_fastapi/extensions/stac_fastapi/extensions/core/pagination/request.py

@@ -34,3 +34,16 @@ class POSTPagination(BaseModel):
     """Page based pagination for POST requests."""
 
     page: Optional[str] = None
+
+
+@attr.s
+class GETOffsetPagination(APIRequest):
+    """Offset pagination for GET requests."""
+
+    offset: Annotated[Optional[int], Query()] = attr.ib(default=None)
+
+
+class POSTOffsetPagination(BaseModel):
+    """Offset pagination model for POST requests."""
+
+    offset: Optional[int] = None

stac_fastapi/extensions/tests/test_pagination.py

@@ -0,0 +1,132 @@
+from typing import Iterator
+
+import pytest
+from starlette.testclient import TestClient
+
+from stac_fastapi.api.app import StacApi
+from stac_fastapi.api.models import (
+    EmptyRequest,
+    create_post_request_model,
+    create_request_model,
+)
+from stac_fastapi.extensions.core import (
+    OffsetPaginationExtension,
+    PaginationExtension,
+    TokenPaginationExtension,
+)
+from stac_fastapi.types.config import ApiSettings
+from stac_fastapi.types.core import BaseCoreClient
+from stac_fastapi.types.search import BaseSearchGetRequest
+
+
+class DummyCoreClient(BaseCoreClient):
+    def all_collections(self, *args, **kwargs):
+        _ = kwargs.pop("request", None)
+        return args, kwargs
+
+    def get_collection(self, *args, **kwargs):
+        _ = kwargs.pop("request", None)
+        return args, kwargs
+
+    def get_item(self, *args, **kwargs):
+        _ = kwargs.pop("request", None)
+        return args, kwargs
+
+    def get_search(self, *args, **kwargs):
+        _ = kwargs.pop("request", None)
+        return args, kwargs
+
+    def post_search(self, *args, **kwargs):
+        _ = kwargs.pop("request", None)
+        return args[0].model_dump(), kwargs
+
+    def item_collection(self, *args, **kwargs):
+        _ = kwargs.pop("request", None)
+        return args, kwargs
+
+
+collections_get_request_model = create_request_model(
+    model_name="CollectionsGetRequest",
+    base_model=EmptyRequest,
+    mixins=[
+        OffsetPaginationExtension().GET,
+    ],
+    request_type="GET",
+)
+
+items_get_request_model = create_request_model(
+    model_name="ItemsGetRequest",
+    base_model=EmptyRequest,
+    mixins=[
+        PaginationExtension().GET,
+    ],
+    request_type="GET",
+)
+
+search_get_request_model = create_request_model(
+    model_name="SearchGetRequest",
+    base_model=BaseSearchGetRequest,
+    mixins=[
+        TokenPaginationExtension().GET,
+    ],
+    request_type="GET",
+)
+
+
+@pytest.fixture
+def client() -> Iterator[TestClient]:
+    settings = ApiSettings()
+
+    api = StacApi(
+        settings=settings,
+        client=DummyCoreClient(),
+        extensions=[],
+        collections_get_request_model=collections_get_request_model,
+        items_get_request_model=items_get_request_model,
+        search_get_request_model=search_get_request_model,
+        search_post_request_model=create_post_request_model([]),
+    )
+    with TestClient(api.app) as client:
+        yield client
+
+
+def test_pagination_extension(client: TestClient):
+    """Test endpoints with pagination extensions."""
+    # OffsetPaginationExtension
+    response = client.get("/collections")
+    assert response.is_success, response.json()
+    arg, kwargs = response.json()
+    assert "offset" in kwargs
+    assert kwargs["offset"] is None
+
+    response = client.get("/collections", params={"offset": 1})
+    assert response.is_success, response.json()
+    arg, kwargs = response.json()
+    assert "offset" in kwargs
+    assert kwargs["offset"] == 1
+
+    # PaginationExtension
+    response = client.get("/collections/a_collection/items")
+    assert response.is_success, response.json()
+    arg, kwargs = response.json()
+    assert "page" in kwargs
+    assert kwargs["page"] is None
+
+    response = client.get("/collections/a_collection/items", params={"page": "1"})
+    assert response.is_success, response.json()
+    arg, kwargs = response.json()
+    assert "page" in kwargs
+    assert kwargs["page"] == "1"
+
+    # TokenPaginationExtension
+    response = client.get("/search")
+    assert response.is_success, response.json()
+    arg, kwargs = response.json()
+    assert "token" in kwargs
+    assert kwargs["token"] is None
+
+    response = client.get("/search", params={"token": "atoken"})
+    assert response.is_success, response.json()
+    arg, kwargs = response.json()
+    assert "token" in kwargs
+    assert kwargs["token"] == "atoken"