feat: paging (#37)

- Closes #30

Author: gadomski

README.md

@@ -31,6 +31,23 @@ docker compose up
 
 This will start **pgstac** on <http://127.0.0.1:8000> and **stac-fastapi** on <http://127.0.0.1:8001>.
 
+### **pgstac**
+
+Right now we don't auto-load data to **pgstac**.
+To populate the **pgstac** database, `uv pip install stacrs-cli` and make sure you've got `docker compose up`, then:
+
+```shell
+curl http://localhost:8001/collections/naip  | pypgstac --dsn postgresql://username:password@localhost:5432/pgstac load collections stdin
+stacrs translate data/naip.parquet | jq -c '.features.[]' | pypgstac --dsn postgresql://username:password@localhost:5432/pgstac load items stdin
+```
+
+To run just the **pgstac** dev server:
+
+```shell
+docker compose up -d pgstac
+TISTAC_BACKEND=postgresql://username:password@localhost:5432/pgstac fastapi dev src/tistac/main.py
+```
+
 ## Developing
 
 Get [yarn](https://yarnpkg.com/getting-started/install), then:

docker-compose.yml

@@ -21,6 +21,8 @@ services:
     build: .
     environment:
       - TISTAC_BACKEND=postgresql://username:password@pgstac:5432/pgstac
+      - TISTAC_CATALOG_ID=tistac-pgstac
+      - TISTAC_CATALOG_DESCRIPTION="tistac with a pgstac backend"
       - UVICORN_HOST=0.0.0.0
       - UVICORN_PORT=8000
     ports:
@@ -33,6 +35,8 @@ services:
     build: .
     environment:
       - TISTAC_BACKEND=/app/data/naip.parquet
+      - TISTAC_CATALOG_ID=tistac-stac-geoparquet
+      - TISTAC_CATALOG_DESCRIPTION="tistac with a stac-geoparquet backend"
       - UVICORN_HOST=0.0.0.0
       - UVICORN_PORT=8001
     ports:

pyproject.toml

@@ -41,6 +41,10 @@ filterwarnings = [
 [tool.ruff.lint]
 select = ["E", "F", "I"]
 
+[tool.uv.sources]
+pgstacrs = { git = "https://github.com/stac-utils/pgstacrs.git" }
+stacrs = { git = "https://github.com/gadomski/stacrs.git" }
+
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"

src/tistac/backends/pgstac.py

@@ -10,8 +10,9 @@ class PgstacBackend(Backend):
     """A PgSTAC backend."""
 
     @classmethod
-    async def open(cls, dsn: str) -> PgstacBackend:
+    async def open(cls, dsn: str, base_url: str) -> PgstacBackend:
         client = await Client.open(dsn)
+        await client.set_setting("base_url", base_url)
         return PgstacBackend(client)
 
     def __init__(self, client: Client) -> None:
@@ -29,5 +30,5 @@ async def get_collection(self, collection_id: str) -> Collection | None:
             return None
 
     async def search(self, search: Search) -> ItemCollection:
-        item_collection = await self.client.search(limit=search.limit)
+        item_collection = await self.client.search(**search.model_dump())
         return ItemCollection.model_validate(item_collection)

src/tistac/backends/stac_geoparquet.py

@@ -1,15 +1,16 @@
 import stacrs
 from pystac import Collection as PystacCollection
 from pystac import Extent, Item
+from stacrs import DuckdbClient
 
 from tistac.backends import Backend
-from tistac.models import Collection, ItemCollection, Search
+from tistac.models import Collection, ItemCollection, Link, Search
 
 
 class StacGeoparquetBackend(Backend):
     """A stac-geoparquet backend, using DuckDB under the hood."""
 
-    def __init__(self, href: str):
+    def __init__(self, href: str, base_url: str):
         # TODO support multiple collections
         # TODO store collection information in the **stac-geoparquet**
         items_as_dicts = stacrs.search(href)
@@ -35,6 +36,8 @@ def __init__(self, href: str):
         d["stac_version"] = "1.1.0"
         self.collections = [Collection.model_validate(d)]
         self.href = href
+        self.base_url = base_url
+        self.client = DuckdbClient()
 
     async def get_collections(self) -> list[Collection]:
         return self.collections
@@ -43,5 +46,27 @@ async def get_collection(self, collection_id: str) -> Collection | None:
         return next((c for c in self.collections if c.id == collection_id), None)
 
     async def search(self, search: Search) -> ItemCollection:
-        items = stacrs.search(self.href, limit=search.limit)
-        return ItemCollection(features=items)
+        item_collection = self.client.search(self.href, **search.model_dump())
+        number_returned = len(item_collection["features"])
+        assert search.limit, "Search should always have a limit at this point"
+        links = []
+        if number_returned >= search.limit:
+            links.append(
+                self.next_link(
+                    search.limit,
+                    search.offset or 0,
+                    number_returned,
+                )
+            )
+        item_collection["numberReturned"] = number_returned
+        item_collection["links"] = links
+        return ItemCollection.model_validate(item_collection)
+
+    def next_link(self, limit: int, offset: int, num_items: int) -> Link:
+        # TODO support POST
+        return Link(
+            rel="next",
+            type="application/geo+json",
+            href=self.base_url + f"search?limit={limit}&offset={offset + num_items}",
+            method="GET",
+        )

src/tistac/dependencies.py

@@ -2,7 +2,7 @@
 from functools import lru_cache
 from typing import Annotated
 
-from fastapi import Depends
+from fastapi import Depends, Request
 
 from tistac.backends import Backend
 from tistac.backends.pgstac import PgstacBackend
@@ -20,19 +20,26 @@ def get_settings() -> Settings:
 # I really don't like how this works -- we might want to go back to leaning on
 # the starlette state, like stac-fastapi-pgstac does. It's just tricky to
 # configure settings w/o using the dependency overrides.
-async def get_backend(settings: Annotated[Settings, Depends(get_settings)]) -> Backend:
+async def get_backend(
+    settings: Annotated[Settings, Depends(get_settings)],
+    request: Request,
+) -> Backend:
     global BACKEND
 
     url = urllib.parse.urlparse(settings.backend)
     if url.scheme == "postgresql":
         if "pgstac" in BACKEND:
             return BACKEND["pgstac"]
         else:
-            BACKEND["pgstac"] = await PgstacBackend.open(settings.backend)
+            BACKEND["pgstac"] = await PgstacBackend.open(
+                settings.backend, str(request.url_for("root"))
+            )
             return BACKEND["pgstac"]
     else:
         if "stac-geoparquet" in BACKEND:
             return BACKEND["stac-geoparquet"]
         else:
-            BACKEND["stac-geoparquet"] = StacGeoparquetBackend(settings.backend)
+            BACKEND["stac-geoparquet"] = StacGeoparquetBackend(
+                settings.backend, str(request.url_for("root"))
+            )
             return BACKEND["stac-geoparquet"]

src/tistac/models/__init__.py

@@ -1,7 +1,6 @@
 from tistac.models.collection import Collection
 from tistac.models.item_collection import ItemCollection
 from tistac.models.link import Link
-from tistac.models.page import Page
 from tistac.models.root import Root
 from tistac.models.search import GetSearch, Search
 
@@ -10,7 +9,6 @@
     "GetSearch",
     "ItemCollection",
     "Link",
-    "Page",
     "Root",
     "Search",
 ]

src/tistac/models/item_collection.py

@@ -1,6 +1,8 @@
-from typing import Any
+from typing import Any, Literal
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
+
+from tistac.models.link import Link
 
 
 class ItemCollection(BaseModel):
@@ -10,4 +12,7 @@ class ItemCollection(BaseModel):
     extension is used.
     """
 
+    type: Literal["FeatureCollection"] = Field(default="FeatureCollection")
+    links: list[Link]
+    number_returned: int = Field(alias="numberReturned")
     features: list[dict[str, Any]]

src/tistac/models/page.py

@@ -1,16 +1,29 @@
-from pydantic import BaseModel, Field
+from __future__ import annotations
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from tistac import Settings
 
 
 class Search(BaseModel):
     """A STAC API POST search."""
 
+    model_config = ConfigDict(extra="allow")
+
     limit: int | None = Field(default=None)
-    """The maximum number of items per page."""
+    offset: int | None = Field(default=None)
+
+    def with_settings(self, settings: Settings) -> Search:
+        if self.limit is None:
+            self.limit = settings.default_limit
+        return self
 
 
 class GetSearch(BaseModel):
     """A STAC API GET search."""
 
+    model_config = ConfigDict(extra="allow")
+
     def into_search(self) -> Search:
         """Converts this GET search into a POST search."""
-        return Search()
+        return Search(**self.model_dump())