Skip to content

Commit

Permalink
feat: allow tokens for paging (#23)
Browse files Browse the repository at this point in the history
  • Loading branch information
gadomski authored Jan 3, 2025
1 parent 7d3ea43 commit 9e1a3b3
Show file tree
Hide file tree
Showing 4 changed files with 37 additions and 5 deletions.
8 changes: 4 additions & 4 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions pgstacrs.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ class Client:
sortby: str | list[str] | None = None,
filter: str | dict[str, Any] | None = None,
query: dict[str, Any] | None = None,
**kwargs: str,
) -> dict[str, Any]:
"""
Searches the database with STAC API item search.
Expand All @@ -59,6 +60,7 @@ class Client:
query: Additional filtering based on properties.
It is recommended to use filter instead, if possible.
limit: The page size returned from the server.
kwargs: Any additional arguments to pass down into the search, e.g a pagination token
"""

async def print_config(self) -> None:
Expand Down
4 changes: 3 additions & 1 deletion src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -293,7 +293,7 @@ impl Client {
})
}

#[pyo3(signature = (*, collections=None, ids=None, intersects=None, bbox=None, datetime=None, include=None, exclude=None, sortby=None, filter=None, query=None, limit=None))]
#[pyo3(signature = (*, collections=None, ids=None, intersects=None, bbox=None, datetime=None, include=None, exclude=None, sortby=None, filter=None, query=None, limit=None, **kwargs))]
fn search<'a>(
&self,
py: Python<'a>,
Expand All @@ -308,6 +308,7 @@ impl Client {
filter: Option<StringOrDict>,
query: Option<Bound<'a, PyDict>>,
limit: Option<u64>,
kwargs: Option<Bound<'a, PyDict>>,
) -> PyResult<Bound<'a, PyAny>> {
let search = stac_api::python::search(
intersects,
Expand All @@ -321,6 +322,7 @@ impl Client {
sortby,
filter,
query,
kwargs,
)?;
self.run(py, |pool| async move {
let connection = pool.get().await?;
Expand Down
28 changes: 28 additions & 0 deletions tests/test_search.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import copy
import urllib.parse
from typing import Any

import pytest
Expand Down Expand Up @@ -219,3 +220,30 @@ async def test_collections(

feature_collection = await client.search(collections=["simple-collection"])
assert len(feature_collection["features"]) == 1


async def test_pagination(
client: Client, collection: dict[str, Any], item: dict[str, Any]
) -> None:
version = await client.get_version()

await client.create_collection(collection)
item["id"] = "first-item"
await client.create_item(item)
second_item = copy.deepcopy(item)
second_item["id"] = "second-item"
await client.create_item(second_item)

feature_collection = await client.search(limit=1, sortby="id")

if version.startswith("0.9"):
next_link = next(
(link for link in feature_collection["links"] if link["rel"] == "next")
)
url = urllib.parse.urlparse(next_link["href"])
token = url.query.split("=")[1]
elif version.startswith("0.8"):
token = "next:" + feature_collection["next"]

feature_collection = await client.search(limit=1, sortby="id", token=token)
assert feature_collection["features"][0]["id"] == "second-item"

0 comments on commit 9e1a3b3

Please sign in to comment.