first pass

Author: jonhealy1

stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/core.py

@@ -4,16 +4,17 @@
 from datetime import datetime as datetime_type
 from datetime import timezone
 from typing import Any, Dict, List, Optional, Set, Type, Union
-from urllib.parse import urljoin
+from base64 import urlsafe_b64encode
 
 import attr
 import stac_pydantic
 from fastapi import HTTPException
 from overrides import overrides
 from pydantic import ValidationError
+from starlette.requests import Request
 from stac_pydantic.links import Relations
 from stac_pydantic.shared import MimeTypes
-from starlette.requests import Request
+from urllib.parse import urljoin
 
 from stac_fastapi.elasticsearch import serializers
 from stac_fastapi.elasticsearch.config import ElasticsearchSettings
@@ -70,7 +71,12 @@ class CoreClient(AsyncBaseCoreClient):
     database = DatabaseLogic()
 
     @overrides
-    async def all_collections(self, **kwargs) -> Collections:
+    async def all_collections(
+        self, 
+        limit: Optional[int] = 10,
+        token: Optional[str] = None,
+        **kwargs
+    ) -> Collections:
         """Read all collections from the database.
 
         Returns:
@@ -80,30 +86,49 @@ async def all_collections(self, **kwargs) -> Collections:
         Raises:
             Exception: If any error occurs while reading the collections from the database.
         """
+        request: Request = kwargs["request"]
         base_url = str(kwargs["request"].base_url)
 
+        hits = self.database.get_all_collections(limit=limit, token=token)
+
+        next_search_after = None
+        next_link = None
+        if len(hits) == limit:
+            last_hit = hits[-1]
+            next_search_after = last_hit['sort']
+            next_token = urlsafe_b64encode(','.join(map(str, next_search_after)).encode()).decode()
+            paging_links = PagingLinks(
+                next=next_token, request=request
+            )
+            next_link = paging_links.link_next()
+
+        links=[
+            {
+                "rel": Relations.root.value,
+                "type": MimeTypes.json,
+                "href": base_url,
+            },
+            {
+                "rel": Relations.parent.value,
+                "type": MimeTypes.json,
+                "href": base_url,
+            },
+            {
+                "rel": Relations.self.value,
+                "type": MimeTypes.json,
+                "href": urljoin(base_url, "collections"),
+            }
+        ]
+
+        if next_link:
+            links.append(next_link)
+        
         return Collections(
             collections=[
-                self.collection_serializer.db_to_stac(c, base_url=base_url)
-                for c in await self.database.get_all_collections()
-            ],
-            links=[
-                {
-                    "rel": Relations.root.value,
-                    "type": MimeTypes.json,
-                    "href": base_url,
-                },
-                {
-                    "rel": Relations.parent.value,
-                    "type": MimeTypes.json,
-                    "href": base_url,
-                },
-                {
-                    "rel": Relations.self.value,
-                    "type": MimeTypes.json,
-                    "href": urljoin(base_url, "collections"),
-                },
+                self.collection_serializer.db_to_stac(c["_source"], base_url=base_url)
+                for c in hits
             ],
+            links=links
         )
 
     @overrides

stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/database_logic.py

@@ -18,6 +18,7 @@
 from stac_fastapi.types.errors import ConflictError, NotFoundError
 from stac_fastapi.types.stac import Collection, Item
 
+
 logger = logging.getLogger(__name__)
 
 NumType = Union[float, int]
@@ -295,9 +296,17 @@ class DatabaseLogic:
 
     """CORE LOGIC"""
 
-    async def get_all_collections(self) -> Iterable[Dict[str, Any]]:
+    async def get_all_collections(
+            self, 
+            token: Optional[str], 
+            limit: int
+    ) -> Iterable[Dict[str, Any]]:
         """Retrieve a list of all collections from the database.
 
+        Args:
+            token (Optional[str]): The token used to return the next set of results.
+            limit (int): Number of results to return 
+
         Returns:
             collections (Iterable[Dict[str, Any]]): A list of dictionaries containing the source data for each collection.
 
@@ -306,10 +315,16 @@ async def get_all_collections(self) -> Iterable[Dict[str, Any]]:
             with the `COLLECTIONS_INDEX` as the target index and `size=1000` to retrieve up to 1000 records.
             The result is a generator of dictionaries containing the source data for each collection.
         """
-        # 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)
-        return (c["_source"] for c in collections["hits"]["hits"])
+        search_after = None
+        if token:
+            search_after = urlsafe_b64decode(token.encode()).decode().split(",")
+        collections = await self.client.search(
+            index=COLLECTIONS_INDEX, 
+            search_after=search_after,
+            size=limit
+        )
+        hits = collections["hits"]["hits"]
+        return hits
 
     async def get_one_item(self, collection_id: str, item_id: str) -> Dict:
         """Retrieve a single item from the database.