diff --git a/CHANGES.md b/CHANGES.md index 226791a1..24e507ac 100644 --- a/CHANGES.md +++ b/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 diff --git a/stac_fastapi/extensions/stac_fastapi/extensions/core/__init__.py b/stac_fastapi/extensions/stac_fastapi/extensions/core/__init__.py index 7fb122e8..77e45d0a 100644 --- a/stac_fastapi/extensions/stac_fastapi/extensions/core/__init__.py +++ b/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", diff --git a/stac_fastapi/extensions/stac_fastapi/extensions/core/pagination/__init__.py b/stac_fastapi/extensions/stac_fastapi/extensions/core/pagination/__init__.py index 0ae847c4..560f0c76 100644 --- a/stac_fastapi/extensions/stac_fastapi/extensions/core/pagination/__init__.py +++ b/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"] diff --git a/stac_fastapi/extensions/stac_fastapi/extensions/core/pagination/offset_pagination.py b/stac_fastapi/extensions/stac_fastapi/extensions/core/pagination/offset_pagination.py new file mode 100644 index 00000000..bc34344f --- /dev/null +++ b/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 diff --git a/stac_fastapi/extensions/stac_fastapi/extensions/core/pagination/request.py b/stac_fastapi/extensions/stac_fastapi/extensions/core/pagination/request.py index 66391c7f..ffa2e322 100644 --- a/stac_fastapi/extensions/stac_fastapi/extensions/core/pagination/request.py +++ b/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 diff --git a/stac_fastapi/extensions/tests/test_pagination.py b/stac_fastapi/extensions/tests/test_pagination.py new file mode 100644 index 00000000..ba1d5b31 --- /dev/null +++ b/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"