paginate, sort collections

Author: jonhealy1

stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/core.py

@@ -1,20 +1,20 @@
 """Item crud client."""
 import json
 import logging
+from base64 import urlsafe_b64encode
 from datetime import datetime as datetime_type
 from datetime import timezone
 from typing import Any, Dict, List, Optional, Set, Type, Union
-from base64 import urlsafe_b64encode
+from urllib.parse import urljoin
 
 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 urllib.parse import urljoin
+from starlette.requests import Request
 
 from stac_fastapi.elasticsearch import serializers
 from stac_fastapi.elasticsearch.config import ElasticsearchSettings
@@ -71,12 +71,7 @@ class CoreClient(AsyncBaseCoreClient):
     database = DatabaseLogic()
 
     @overrides
-    async def all_collections(
-        self, 
-        limit: Optional[int] = 10,
-        token: Optional[str] = None,
-        **kwargs
-    ) -> Collections:
+    async def all_collections(self, **kwargs) -> Collections:
         """Read all collections from the database.
 
         Returns:
@@ -89,20 +84,23 @@ async def all_collections(
         request: Request = kwargs["request"]
         base_url = str(kwargs["request"].base_url)
 
-        hits = self.database.get_all_collections(limit=limit, token=token)
+        limit = int(request.query_params["limit"]) if "limit" in request.query_params else 10
+        token = request.query_params["token"] if "token" in request.query_params else None
+
+        hits = await 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_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=[
+        links = [
             {
                 "rel": Relations.root.value,
                 "type": MimeTypes.json,
@@ -117,18 +115,18 @@ async def all_collections(
                 "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["_source"], base_url=base_url)
                 for c in hits
             ],
-            links=links
+            links=links,
         )
 
     @overrides

stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/database_logic.py

@@ -18,7 +18,6 @@
 from stac_fastapi.types.errors import ConflictError, NotFoundError
 from stac_fastapi.types.stac import Collection, Item
 
-
 logger = logging.getLogger(__name__)
 
 NumType = Union[float, int]
@@ -297,31 +296,27 @@ class DatabaseLogic:
     """CORE LOGIC"""
 
     async def get_all_collections(
-            self, 
-            token: Optional[str], 
-            limit: int
+        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 
+            limit (int): Number of results to return
 
         Returns:
             collections (Iterable[Dict[str, Any]]): A list of dictionaries containing the source data for each collection.
 
         Notes:
             The collections are retrieved from the Elasticsearch database using the `client.search` method,
-            with the `COLLECTIONS_INDEX` as the target index and `size=1000` to retrieve up to 1000 records.
+            with the `COLLECTIONS_INDEX` as the target index and `size=limit` to retrieve records.
             The result is a generator of dictionaries containing the source data for each collection.
         """
         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
+            index=COLLECTIONS_INDEX, search_after=search_after, size=limit, sort={"id": {"order": "asc"}}
         )
         hits = collections["hits"]["hits"]
         return hits