Add pagination (#9)

* bulk ingest test data
* use pgstac 0.9.0
* add tests for pagination
* add pagination

Author: ungarj

docker-compose.yml

@@ -29,7 +29,7 @@ services:
     command: bash -c "./scripts/wait-for-it.sh database:5432 && python -m stac_fastapi_pgstac_pair_search.app"
 
   database:
-    image: ghcr.io/stac-utils/pgstac:v0.9.2
+    image: ghcr.io/stac-utils/pgstac:v0.9.8
     environment:
       - POSTGRES_USER=username
       - POSTGRES_PASSWORD=password

scripts/ingest_test_data.py

@@ -89,14 +89,21 @@ def ingest_data(
             post_or_put(urljoin(app_host, "/collections"), json.loads(src.read()))
 
         # ingest downloaded items
-        for item_json in tqdm.tqdm(
-            os.listdir(collection_dir + "/items"), desc="ingesting items", unit=" items"
-        ):
-            with open(f"{collection_dir}/items/{item_json}", "r") as src:
-                post_or_put(
-                    urljoin(app_host, f"/collections/{collection_name}/items"),
-                    json.loads(src.read()),
+        item_collection = {
+            "type": "FeatureCollection",
+            "features": [
+                json.loads(open(f"{collection_dir}/items/{item_json}").read())
+                for item_json in tqdm.tqdm(
+                    os.listdir(collection_dir + "/items"),
+                    desc="ingesting items",
+                    unit=" items",
                 )
+            ],
+        }
+        post_or_put(
+            urljoin(app_host, f"/collections/{collection_name}/items"),
+            item_collection,
+        )
 
     # load queryables
     print("Loading queryables")

src/stac_fastapi_pgstac_pair_search/client.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+import json
 import logging
 from pathlib import Path
 from typing import Dict, Any, Optional, Set, List, Tuple
@@ -12,7 +13,7 @@
 from stac_fastapi.pgstac.app import StacApi
 from stac_fastapi.pgstac.config import Settings
 from stac_fastapi.pgstac.core import CoreCrudClient
-from stac_fastapi.pgstac.models.links import ItemLinks, PagingLinks
+from stac_fastapi.pgstac.models.links import ItemLinks
 from stac_fastapi.pgstac.utils import filter_fields
 from stac_fastapi.types.errors import InvalidQueryParameter
 from stac_fastapi.types.stac import ItemCollection
@@ -122,17 +123,6 @@ async def _pair_search_base(
             raise InvalidQueryParameter(
                 f"Datetime parameter {search_request.datetime} is invalid."
             ) from e
-        # Starting in pgstac 0.9.0, the `next` and `prev` tokens are returned in spec-compliant links with method GET
-        next_from_link: Optional[str] = None
-        prev_from_link: Optional[str] = None
-        for link in items.get("links", []):
-            if link.get("rel") == "next":
-                next_from_link = link.get("href").split("token=next:")[1]
-            if link.get("rel") == "prev":
-                prev_from_link = link.get("href").split("token=prev:")[1]
-
-        next: Optional[str] = items.pop("next", next_from_link)
-        prev: Optional[str] = items.pop("prev", prev_from_link)
         collection = ItemCollection(**items)
 
         fields = getattr(search_request, "fields", None)
@@ -191,13 +181,70 @@ async def _get_base_item(collection_id: str) -> Dict[str, Any]:
                 cleaned_features.append(feature)
 
         collection["features"] = cleaned_features
-        collection["links"] = await PagingLinks(
-            request=request,
-            next=next,
-            prev=prev,
-        ).get_links()
+        if cleaned_features:
+            collection["links"] = await self._get_search_links(
+                search_request=search_request, request=request
+            )
+
         return collection
 
+    async def _get_search_links(
+        self, search_request: PairSearchRequest, request: Request
+    ) -> List[Dict[str, str]]:
+        """Take existing request and edit offset."""
+        links = []
+        next_page_offset = search_request.offset + search_request.limit
+        prev_page_offset = max(search_request.offset - search_request.limit, 0)
+        if request.method == "GET":
+            query_params = request.query_params.multi_items()
+            links.append(
+                {
+                    "rel": "next",
+                    "href": request.url.replace_query_params(
+                        **dict(query_params, offset=next_page_offset)
+                    ),
+                    "type": "application/geo+json",
+                }
+            )
+            # add link to previous page
+            if search_request.offset:
+                links.append(
+                    {
+                        "rel": "prev",
+                        "href": request.url.replace_query_params(
+                            **dict(
+                                query_params,
+                                offset=prev_page_offset,
+                            )
+                        ),
+                        "type": "application/geo+json",
+                    }
+                )
+        elif request.method == "POST":
+            body = await request.body()
+            links.append(
+                {
+                    "rel": "next",
+                    "href": request.url,
+                    "type": "application/geo+json",
+                    "method": "POST",
+                    "body": dict(json.loads(body), offset=next_page_offset),
+                }
+            )
+            # add link to previous page
+            if search_request.offset:
+                links.append(
+                    {
+                        "rel": "prev",
+                        "href": request.url,
+                        "type": "application/geo+json",
+                        "method": "POST",
+                        "body": dict(json.loads(body), offset=prev_page_offset),
+                    }
+                )
+
+        return links
+
 
 def render_sql(pair_search_request: PairSearchRequest) -> Tuple[str, List[Any]]:
     filter_query, filter_params = pair_search_request.filter_sql
@@ -212,6 +259,7 @@ def render_sql(pair_search_request: PairSearchRequest) -> Tuple[str, List[Any]]:
             exclude_none=True, by_alias=True
         ),
         limit=pair_search_request.limit or 10,
+        offset=pair_search_request.offset or 0,
         response_type=pair_search_request.response_type,
         **filter_params,
     )

src/stac_fastapi_pgstac_pair_search/models.py

@@ -4,7 +4,7 @@
 
 import cql2
 from fastapi import Query
-from pydantic import Field, AfterValidator, BaseModel, model_validator
+from pydantic import Field, AfterValidator, BaseModel, model_validator, NonNegativeInt
 from datetime import datetime as dt
 from stac_fastapi.extensions.core.filter.request import FilterLang
 from stac_fastapi.types.search import Limit, APIRequest, BaseSearchPostRequest
@@ -42,6 +42,11 @@ class PairSearchRequest(BaseModel, APIRequest):
         description="Limits the number of results that are included in each page of the response (capped to 10_000).",  # noqa: E501
     )
 
+    offset: Optional[NonNegativeInt] = Field(
+        0,
+        description="Offset from the first record. Can be used to paginate results.",
+    )
+
     first_bbox: Annotated[Optional[BBox], AfterValidator(validate_bbox)] = Field(
         alias="first-bbox", default=None
     )

src/stac_fastapi_pgstac_pair_search/sql/pair_search.sql

@@ -17,7 +17,7 @@ all_pairs AS (
 limited_pairs AS (
     SELECT id1, id2, first, second
     FROM all_pairs
-    OFFSET 0
+    OFFSET :offset
     LIMIT :limit
 ),
 all_features AS (