Skip to content

Commit

Permalink
update pagination for collection-search (#155)
Browse files Browse the repository at this point in the history
* add failing test

* handle collection paging differently (#156)

* handle collection paging differently

* test next link

* add OffsetPaginationExtension

* uncomment test

* update to pgstac 0.9.2

* update prev logic

* only test 0.9.0

* update pypstac version

* add back 0.8 support but allow skip tests

* skip for 0.8

* remove warnings

* fallback to all_collections when `CollectionSearchExtension` is not enabled (#179)

* fallback to all_collections when `CollectionSearchExtension` is not enabled

* test all_collection fallback

* add offset=0

* Update tests/conftest.py

Co-authored-by: Henry Rodman <henry.rodman@gmail.com>

---------

Co-authored-by: Henry Rodman <henry.rodman@gmail.com>
  • Loading branch information
vincentsarago and hrodmn authored Jan 13, 2025
1 parent 7135059 commit 0178a4f
Show file tree
Hide file tree
Showing 8 changed files with 361 additions and 37 deletions.
1 change: 1 addition & 0 deletions .github/workflows/cicd.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ jobs:
- {python: '3.12', pypgstac: '0.9.*'}
- {python: '3.12', pypgstac: '0.8.*'}
- {python: '3.11', pypgstac: '0.8.*'}
- {python: '3.10', pypgstac: '0.8.*'}
- {python: '3.9', pypgstac: '0.8.*'}
- {python: '3.8', pypgstac: '0.8.*'}

Expand Down
4 changes: 3 additions & 1 deletion docker-compose.yml
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ services:
build:
context: .
dockerfile: Dockerfile.tests
volumes:
- .:/app
environment:
- ENVIRONMENT=local
- DB_MIN_CONN_SIZE=1
Expand All @@ -44,7 +46,7 @@ services:

database:
container_name: stac-db
image: ghcr.io/stac-utils/pgstac:v0.9.1
image: ghcr.io/stac-utils/pgstac:v0.9.2
environment:
- POSTGRES_USER=username
- POSTGRES_PASSWORD=password
Expand Down
6 changes: 3 additions & 3 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,9 @@
"orjson",
"pydantic",
"stac_pydantic==3.1.*",
"stac-fastapi.api~=3.0.2",
"stac-fastapi.extensions~=3.0.2",
"stac-fastapi.types~=3.0.2",
"stac-fastapi.api~=3.0.3",
"stac-fastapi.extensions~=3.0.3",
"stac-fastapi.types~=3.0.3",
"asyncpg",
"buildpg",
"brotli_asgi",
Expand Down
2 changes: 2 additions & 0 deletions stac_fastapi/pgstac/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
from stac_fastapi.extensions.core import (
FieldsExtension,
FilterExtension,
OffsetPaginationExtension,
SortExtension,
TokenPaginationExtension,
TransactionExtension,
Expand Down Expand Up @@ -58,6 +59,7 @@
"sort": SortExtension(),
"fields": FieldsExtension(),
"filter": FilterExtension(client=FiltersClient()),
"pagination": OffsetPaginationExtension(),
}

enabled_extensions = (
Expand Down
80 changes: 47 additions & 33 deletions stac_fastapi/pgstac/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
from stac_fastapi.pgstac.config import Settings
from stac_fastapi.pgstac.models.links import (
CollectionLinks,
CollectionSearchPagingLinks,
ItemCollectionLinks,
ItemLinks,
PagingLinks,
Expand All @@ -46,8 +47,8 @@ async def all_collections( # noqa: C901
bbox: Optional[BBox] = None,
datetime: Optional[DateTimeType] = None,
limit: Optional[int] = None,
offset: Optional[int] = None,
query: Optional[str] = None,
token: Optional[str] = None,
fields: Optional[List[str]] = None,
sortby: Optional[str] = None,
filter: Optional[str] = None,
Expand All @@ -64,38 +65,51 @@ async def all_collections( # noqa: C901
"""
base_url = get_base_url(request)

# Parse request parameters
base_args = {
"bbox": bbox,
"limit": limit,
"token": token,
"query": orjson.loads(unquote_plus(query)) if query else query,
}

clean_args = clean_search_args(
base_args=base_args,
datetime=datetime,
fields=fields,
sortby=sortby,
filter_query=filter,
filter_lang=filter_lang,
)

async with request.app.state.get_connection(request, "r") as conn:
q, p = render(
"""
SELECT * FROM collection_search(:req::text::jsonb);
""",
req=json.dumps(clean_args),
next_link: Optional[Dict[str, Any]] = None
prev_link: Optional[Dict[str, Any]] = None
collections_result: Collections

if self.extension_is_enabled("CollectionSearchExtension"):
base_args = {
"bbox": bbox,
"limit": limit,
"offset": offset,
"query": orjson.loads(unquote_plus(query)) if query else query,
}

clean_args = clean_search_args(
base_args=base_args,
datetime=datetime,
fields=fields,
sortby=sortby,
filter_query=filter,
filter_lang=filter_lang,
)
collections_result: Collections = await conn.fetchval(q, *p)

next: Optional[str] = None
prev: Optional[str] = None
async with request.app.state.get_connection(request, "r") as conn:
q, p = render(
"""
SELECT * FROM collection_search(:req::text::jsonb);
""",
req=json.dumps(clean_args),
)
collections_result = await conn.fetchval(q, *p)

if links := collections_result.get("links"):
next = collections_result["links"].pop("next")
prev = collections_result["links"].pop("prev")
if links := collections_result.get("links"):
for link in links:
if link["rel"] == "next":
next_link = link
elif link["rel"] == "prev":
prev_link = link

else:
async with request.app.state.get_connection(request, "r") as conn:
cols = await conn.fetchval(
"""
SELECT * FROM all_collections();
"""
)
collections_result = {"collections": cols, "links": []}

linked_collections: List[Collection] = []
collections = collections_result["collections"]
Expand All @@ -120,10 +134,10 @@ async def all_collections( # noqa: C901

linked_collections.append(coll)

links = await PagingLinks(
links = await CollectionSearchPagingLinks(
request=request,
next=next,
prev=prev,
next=next_link,
prev=prev_link,
).get_links()

return Collections(
Expand Down
49 changes: 49 additions & 0 deletions stac_fastapi/pgstac/models/links.py
Original file line number Diff line number Diff line change
Expand Up @@ -173,6 +173,55 @@ def link_prev(self) -> Optional[Dict[str, Any]]:
return None


@attr.s
class CollectionSearchPagingLinks(BaseLinks):
next: Optional[Dict[str, Any]] = attr.ib(kw_only=True, default=None)
prev: Optional[Dict[str, Any]] = attr.ib(kw_only=True, default=None)

def link_next(self) -> Optional[Dict[str, Any]]:
"""Create link for next page."""
if self.next is not None:
method = self.request.method
if method == "GET":
# if offset is equal to default value (0), drop it
if self.next["body"].get("offset", -1) == 0:
_ = self.next["body"].pop("offset")

href = merge_params(self.url, self.next["body"])

# if next link is equal to this link, skip it
if href == self.url:
return None

return {
"rel": Relations.next.value,
"type": MimeTypes.geojson.value,
"method": method,
"href": href,
}

return None

def link_prev(self):
if self.prev is not None:
method = self.request.method
if method == "GET":
href = merge_params(self.url, self.prev["body"])

# if prev link is equal to this link, skip it
if href == self.url:
return None

return {
"rel": Relations.previous.value,
"type": MimeTypes.geojson.value,
"method": method,
"href": href,
}

return None


@attr.s
class CollectionLinksBase(BaseLinks):
"""Create inferred links specific to collections."""
Expand Down
54 changes: 54 additions & 0 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from fastapi import APIRouter
from fastapi.responses import ORJSONResponse
from httpx import ASGITransport, AsyncClient
from pypgstac import __version__ as pgstac_version
from pypgstac.db import PgstacDB
from pypgstac.migrate import Migrate
from pytest_postgresql.janitor import DatabaseJanitor
Expand All @@ -26,6 +27,7 @@
CollectionSearchExtension,
FieldsExtension,
FilterExtension,
OffsetPaginationExtension,
SortExtension,
TokenPaginationExtension,
TransactionExtension,
Expand All @@ -47,6 +49,12 @@
logger = logging.getLogger(__name__)


requires_pgstac_0_9_2 = pytest.mark.skipif(
tuple(map(int, pgstac_version.split("."))) < (0, 9, 2),
reason="PgSTAC>=0.9.2 required",
)


@pytest.fixture(scope="session")
def event_loop():
return asyncio.get_event_loop()
Expand Down Expand Up @@ -140,6 +148,7 @@ def api_client(request, database):
SortExtension(),
FieldsExtension(),
FilterExtension(client=FiltersClient()),
OffsetPaginationExtension(),
]
collection_search_extension = CollectionSearchExtension.from_extensions(
collection_extensions
Expand Down Expand Up @@ -259,3 +268,48 @@ async def load_test2_item(app_client, load_test_data, load_test2_collection):
)
assert resp.status_code == 201
return Item.model_validate(resp.json())


@pytest.fixture(
scope="session",
)
def api_client_no_ext(database):
api_settings = Settings(
postgres_user=database.user,
postgres_pass=database.password,
postgres_host_reader=database.host,
postgres_host_writer=database.host,
postgres_port=database.port,
postgres_dbname=database.dbname,
testing=True,
)
return StacApi(
settings=api_settings,
extensions=[
TransactionExtension(client=TransactionsClient(), settings=api_settings)
],
client=CoreCrudClient(),
)


@pytest.fixture(scope="function")
async def app_no_ext(api_client_no_ext):
logger.info("Creating app Fixture")
time.time()
app = api_client_no_ext.app
await connect_to_db(app)

yield app

await close_db_connection(app)

logger.info("Closed Pools.")


@pytest.fixture(scope="function")
async def app_client_no_ext(app_no_ext):
logger.info("creating app_client")
async with AsyncClient(
transport=ASGITransport(app=app_no_ext), base_url="http://test"
) as c:
yield c
Loading

0 comments on commit 0178a4f

Please sign in to comment.