diff --git a/stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/core.py b/stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/core.py index 2f6d3599..7a3bf1b4 100644 --- a/stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/core.py +++ b/stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/core.py @@ -57,11 +57,20 @@ class CoreClient(AsyncBaseCoreClient): async def all_collections(self, **kwargs) -> Collections: """Read all collections from the database.""" base_url = str(kwargs["request"].base_url) + request = kwargs["request"] + page_str = None + if request.query_params is not None: + page_str = request.query_params.get("page") + pageint = 1 + if page_str is not None: + pageint = int(page_str) + if pageint < 1: + pageint = 1 return Collections( collections=[ self.collection_serializer.db_to_stac(c, base_url=base_url) - for c in await self.database.get_all_collections() + for c in await self.database.get_all_collections(page=pageint) ], links=[ { diff --git a/stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/database_logic.py b/stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/database_logic.py index 568dffc1..1fbfa673 100644 --- a/stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/database_logic.py +++ b/stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/database_logic.py @@ -22,6 +22,7 @@ NumType = Union[float, int] COLLECTIONS_INDEX = "collections" +COLLECTIONS_PAGE_SIZE = 1000 ITEMS_INDEX_PREFIX = "items_" DEFAULT_INDICES = f"*,-*kibana*,-{COLLECTIONS_INDEX}" @@ -210,11 +211,12 @@ class DatabaseLogic: """CORE LOGIC""" - async def get_all_collections(self) -> Iterable[Dict[str, Any]]: + async def get_all_collections(self, page: int = 1) -> Iterable[Dict[str, Any]]: """Database logic to retrieve a list of all collections.""" - # https://github.com/stac-utils/stac-fastapi-elasticsearch/issues/65 - # collections should be paginated, but at least return more than the default 10 for now - collections = await self.client.search(index=COLLECTIONS_INDEX, size=1000) + results_after = (page - 1) * COLLECTIONS_PAGE_SIZE + collections = await self.client.search( + index=COLLECTIONS_INDEX, size=COLLECTIONS_PAGE_SIZE, from_=results_after + ) return (c["_source"] for c in collections["hits"]["hits"]) async def get_one_item(self, collection_id: str, item_id: str) -> Dict: @@ -428,7 +430,7 @@ def sync_prep_create_item(self, item: Item, base_url: str) -> Item: return self.item_serializer.stac_to_db(item, base_url) - async def create_item(self, item: Item, refresh: bool = False): + async def create_item(self, item: Item, refresh: bool = True): """Database logic for creating one item.""" # todo: check if collection exists, but cache item_id = item["id"] @@ -445,9 +447,7 @@ async def create_item(self, item: Item, refresh: bool = False): f"Item {item_id} in collection {collection_id} already exists" ) - async def delete_item( - self, item_id: str, collection_id: str, refresh: bool = False - ): + async def delete_item(self, item_id: str, collection_id: str, refresh: bool = True): """Database logic for deleting one item.""" try: await self.client.delete( @@ -460,7 +460,7 @@ async def delete_item( f"Item {item_id} in collection {collection_id} not found" ) - async def create_collection(self, collection: Collection, refresh: bool = False): + async def create_collection(self, collection: Collection, refresh: bool = True): """Database logic for creating one collection.""" collection_id = collection["id"] @@ -487,7 +487,7 @@ async def find_collection(self, collection_id: str) -> Collection: return collection["_source"] - async def delete_collection(self, collection_id: str, refresh: bool = False): + async def delete_collection(self, collection_id: str, refresh: bool = True): """Database logic for deleting one collection.""" await self.find_collection(collection_id=collection_id) await self.client.delete( diff --git a/stac_fastapi/elasticsearch/tests/conftest.py b/stac_fastapi/elasticsearch/tests/conftest.py index 25bf6850..a9852da6 100644 --- a/stac_fastapi/elasticsearch/tests/conftest.py +++ b/stac_fastapi/elasticsearch/tests/conftest.py @@ -40,6 +40,7 @@ def __init__(self, item, collection): class MockRequest: base_url = "http://test-server" + query_params = None def __init__( self, method: str = "GET", url: str = "XXXX", app: Optional[Any] = None diff --git a/stac_fastapi/elasticsearch/tests/resources/test_collection.py b/stac_fastapi/elasticsearch/tests/resources/test_collection.py index 5172c2b0..f69aec0b 100644 --- a/stac_fastapi/elasticsearch/tests/resources/test_collection.py +++ b/stac_fastapi/elasticsearch/tests/resources/test_collection.py @@ -1,4 +1,5 @@ import pystac +import pytest async def test_create_and_delete_collection(app_client, load_test_data): @@ -13,6 +14,33 @@ async def test_create_and_delete_collection(app_client, load_test_data): assert resp.status_code == 200 +@pytest.mark.skip( + reason="paginating collections takes a long time to test, skip it if you haven't changed anything" +) +async def test_create_paginate_collections(app_client, load_test_data): + """Test creation and pagination of collections""" + test_collection = load_test_data("test_collection.json") + test_collection["id"] = "test" + + for i in range(1, 1005): + test_collection["id"] = "test_" + str(i) + resp = await app_client.post("/collections", json=test_collection) + assert resp.status_code == 200 + + resp = await app_client.get("/collections?page=2") # , params={"page": 2}) + resp_json = resp.json() + collcount = len(resp_json["collections"]) + assert collcount == 4 + + """ this many deletes all at once tends to error out after a certain point """ + + +# for i in range(1, 1005): +# test_id = "test_" + str(i) +# resp = await app_client.delete(f"/collections/{test_id}") +# assert resp.status_code == 200 + + async def test_create_collection_conflict(app_client, ctx): """Test creation of a collection which already exists""" # This collection ID is created in the fixture, so this should be a conflict