Merge pull request #23 from EO-DataHub/bugfix/EODHP-549-fastapi-ensure-next-link-only-included-when-subsequent-results-exist

Bugfix/eodhp 549 fastapi ensure next link only included when subsequent results exist

Author: tjellicoe-tpzuk

stac_fastapi/core/stac_fastapi/core/core.py

@@ -273,19 +273,26 @@ async def all_collections(self, username_header: dict, **kwargs) -> Collections:
         collections = []
 
         while True:
-            temp_collections, next_token = await self.database.get_all_collections(
-                token=token, limit=limit, base_url=base_url
+            temp_collections, next_token, hit_tokens = (
+                await self.database.get_all_collections(
+                    token=token, limit=limit, base_url=base_url
+                )
             )
 
-            for collection in temp_collections:
+            for i, (collection, hit_token) in enumerate(
+                zip(temp_collections, hit_tokens)
+            ):
                 # Get access control array for each collection
                 access_control = collection["access_control"]
                 collection.pop("access_control")
                 # Append collection to list if user has access
                 if int(access_control[-1]) or int(access_control[user_index]):
                     collections.append(collection)
                     if len(collections) >= limit:
-                        break
+                        # Extract token from last result
+                        if i < len(temp_collections) - 1:
+                            next_token = hit_token
+                            break
 
             # If collections now less than limit and more results, will need to run search again, giving next_token
             if len(collections) >= limit or not next_token:
@@ -338,29 +345,33 @@ async def all_catalogs(
 
         while True:
             # Search is run continually until limit is reached or no more results
-            temp_catalogs, next_token = await self.database.get_all_catalogs(
-                catalog_path=catalog_path,
-                token=token,
-                limit=limit,
-                base_url=base_url,
-                user_index=user_index,
-                conformance_classes=self.conformance_classes(),
+            temp_catalogs, next_token, hit_tokens = (
+                await self.database.get_all_catalogs(
+                    catalog_path=catalog_path,
+                    token=token,
+                    limit=limit,
+                    base_url=base_url,
+                    user_index=user_index,
+                    conformance_classes=self.conformance_classes(),
+                )
             )
 
-            for catalog in temp_catalogs:
+            for i, (catalog, hit_token) in enumerate(zip(temp_catalogs, hit_tokens)):
                 # Get access control array for each catalog
                 access_control = catalog["access_control"]
                 catalog.pop("access_control")
                 # Add catalog to list if user has access
                 if int(access_control[-1]) or int(access_control[user_index]):
                     catalogs.append(catalog)
                     if len(catalogs) >= limit:
-                        break
+                        if i < len(temp_catalogs) - 1:
+                            # Extract token from last result
+                            next_token = hit_token
+                            break
 
             # If catalogs now less than limit and more results, will need to run search again, giving next_token
             if len(catalogs) >= limit or not next_token:
                 # TODO: implement smarter token logic to return token of last returned ES entry
-                next_token = token
                 break
             token = next_token
 
@@ -465,7 +476,7 @@ async def get_catalog(
             )
 
         # Assume at most 100 collections in a catalog for the time being, may need to increase
-        collections, _ = await self.database.get_catalog_collections(
+        collections, _, _ = await self.database.get_catalog_collections(
             catalog_path=catalog_path,
             base_url=base_url,
             limit=NUMBER_OF_CATALOG_COLLECTIONS,
@@ -560,28 +571,34 @@ async def get_catalog_collections(
         collections = []
 
         while True:
-            temp_collections, next_token = await self.database.get_catalog_collections(
-                catalog_path=catalog_path,
-                token=token,  # type: ignore
-                limit=limit,
-                base_url=base_url,
+            temp_collections, next_token, hit_tokens = (
+                await self.database.get_catalog_collections(
+                    catalog_path=catalog_path,
+                    token=token,  # type: ignore
+                    limit=limit,
+                    base_url=base_url,
+                )
             )
 
             # Check if current user has access to each collection
-            for collection in temp_collections:
+            for i, (collection, hit_token) in enumerate(
+                zip(temp_collections, hit_tokens)
+            ):
                 # Get access control array for each collection
                 access_control = collection["access_control"]
                 collection.pop("access_control")
                 # Remove collection from list if user does not have access
                 if int(access_control[-1]) or int(access_control[user_index]):
                     collections.append(collection)
                     if len(collections) >= limit:
-                        break
+                        if i < len(temp_collections) - 1:
+                            # Extract token from last result
+                            next_token = hit_token
+                            break
 
             # If collections now less than limit and more results, will need to run search again, giving next_token
             if len(collections) >= limit or not next_token:
                 # TODO: implement smarter token logic to return token of last returned ES entry
-                next_token = token
                 break
             token = next_token
 
@@ -674,7 +691,8 @@ async def item_collection(
 
             search = self.database.apply_bbox_filter(search=search, bbox=bbox)
 
-        items, maybe_count, next_token = await self.database.execute_search(
+        # No further access control needed as already checked above for collection
+        items, maybe_count, next_token, _ = await self.database.execute_search(
             search=search,
             catalog_paths=[catalog_path],
             limit=limit,
@@ -1046,17 +1064,19 @@ async def post_global_search(
         items = []
 
         while True:
-            temp_items, maybe_count, next_token = await self.database.execute_search(
-                search=search,
-                limit=limit,
-                token=token,  # type: ignore
-                sort=sort,
-                collection_ids=search_request.collections,
-                catalog_paths=search_request.catalog_paths,
+            temp_items, maybe_count, next_token, hit_tokens = (
+                await self.database.execute_search(
+                    search=search,
+                    limit=limit,
+                    token=token,  # type: ignore
+                    sort=sort,
+                    collection_ids=search_request.collections,
+                    catalog_paths=search_request.catalog_paths,
+                )
             )
 
             # Filter results to those that are accessible to the user
-            for item in temp_items:
+            for i, (item, hit_token) in enumerate(zip(temp_items, hit_tokens)):
                 # Get item index for path extraction
                 item_catalog_path = item[1]
                 # Get parent collection if collection is present
@@ -1073,7 +1093,10 @@ async def post_global_search(
                     if int(access_control[-1]) or int(access_control[user_index]):
                         items.append(item)
                         if len(items) >= limit:
-                            break
+                            if i < len(temp_items) - 1:
+                                # Extract token from last result
+                                next_token = hit_token
+                                break
                 # Get parent catalog if collection is not present
                 else:
                     # Get access control array for this catalog
@@ -1085,12 +1108,14 @@ async def post_global_search(
                     if int(access_control[-1]) or int(access_control[user_index]):
                         items.append(item)
                         if len(items) >= limit:
-                            break
+                            # Extract token from last result
+                            if i < len(temp_items) - 1:
+                                next_token = hit_token
+                                break
 
             # If items now less than limit and more results, will need to run search again, giving next_token
             if len(items) >= limit or not next_token:
                 # TODO: implement smarter token logic to return token of last returned ES entry
-                next_token = token
                 break
             token = next_token
 
@@ -1264,7 +1289,6 @@ async def post_search(
         username_header: dict,
         **kwargs,
     ) -> ItemCollection:
-        print("post search")
         """
         Perform a POST search on a specific sub-catalog.
 
@@ -1299,9 +1323,12 @@ async def post_search(
             raise HTTPException(
                 status_code=403, detail="User does not have access to this Catalog"
             )
+        collections = []
+        if search_request.collections:
+            collections = search_request.collections
 
         # Filter the search collections to those that are accessible to the user
-        for collection_id in search_request.collections[:]:
+        for collection_id in collections[:]:
             # Filter the search catalogs to those that are accessible to the user
             collection = await self.database.find_collection(
                 catalog_path=catalog_path, collection_id=collection_id
@@ -1310,7 +1337,7 @@ async def post_search(
             access_control = collection["access_control"]
             # Remove catalog from list if user does not have access
             if not int(access_control[-1]) and not int(access_control[user_index]):
-                search_request.collections.remove(collection_id)
+                collections.remove(collection_id)
 
         search = self.database.make_search()
 
@@ -1319,9 +1346,9 @@ async def post_search(
                 search=search, item_ids=search_request.ids
             )
 
-        if search_request.collections:
+        if collections:
             search = self.database.apply_collections_filter(
-                search=search, collection_ids=search_request.collections
+                search=search, collection_ids=collections
             )
 
         if search_request.datetime:
@@ -1375,17 +1402,19 @@ async def post_search(
         items = []
 
         while True:
-            temp_items, maybe_count, next_token = await self.database.execute_search(
-                search=search,
-                limit=limit,
-                token=token,  # type: ignore
-                sort=sort,
-                collection_ids=search_request.collections,
-                catalog_paths=[catalog_path],
+            temp_items, maybe_count, next_token, hit_tokens = (
+                await self.database.execute_search(
+                    search=search,
+                    limit=limit,
+                    token=token,  # type: ignore
+                    sort=sort,
+                    collection_ids=collections,
+                    catalog_paths=[catalog_path],
+                )
             )
 
             # Filter results to those that are accessible to the user
-            for item in temp_items:
+            for i, (item, hit_token) in enumerate(zip(temp_items, hit_tokens)):
                 # Get item index for path extraction
                 item_catalog_path = item[1]
                 # Get parent collection if collection is present
@@ -1401,6 +1430,11 @@ async def post_search(
                     # Append item to list if user has access
                     if int(access_control[-1]) or int(access_control[user_index]):
                         items.append(item)
+                        if len(items) >= limit:
+                            if i < len(temp_items) - 1:
+                                # Extract token from last result
+                                next_token = hit_token
+                                break
                 # Get parent catalog if collection is not present
                 else:
                     # Get access control array for this catalog
@@ -1411,11 +1445,15 @@ async def post_search(
                     # Append item to list if user has access
                     if int(access_control[-1]) or int(access_control[user_index]):
                         items.append(item)
+                        if len(items) >= limit:
+                            if i < len(temp_items) - 1:
+                                # Extract token from last result
+                                next_token = hit_token
+                                break
 
             # If items now less than limit and more results, will need to run search again, giving next_token
             if len(items) >= limit or not next_token:
                 # TODO: implement smarter token logic to return token of last returned ES entry
-                next_token = token
                 break
             token = next_token
 
@@ -2254,7 +2292,7 @@ async def post_all_collections(
         collections = []
 
         while True:
-            temp_collections, _, next_token = (
+            temp_collections, _, next_token, hit_tokens = (
                 await self.database.execute_collection_search(
                     search=search,
                     limit=limit,
@@ -2265,18 +2303,24 @@ async def post_all_collections(
             )
 
             # Filter results to those that are accessible to the user
-            for collection in temp_collections:
+            for i, (collection, hit_token) in enumerate(
+                zip(temp_collections, hit_tokens)
+            ):
                 # Get access control array for this collection
                 access_control = collection["access_control"]
                 collection.pop("access_control")
                 # Append collection to list if user has access
                 if int(access_control[-1]) or int(access_control[user_index]):
                     collections.append(collection)
+                    if len(collections) >= limit:
+                        if i < len(temp_collections) - 1:
+                            # Extract token from last result
+                            next_token = hit_token
+                            break
 
             # If collections now less than limit and more results, will need to run search again, giving next_token
             if len(collections) >= limit or not next_token:
                 # TODO: implement smarter token logic to return token of last returned ES entry
-                next_token = token
                 break
             token = next_token
 
@@ -2411,7 +2455,7 @@ async def post_discovery_search(
         catalogs_and_collections = []
 
         while True:
-            temp_catalogs_and_collections, maybe_count, next_token = (
+            temp_catalogs_and_collections, _, next_token, hit_tokens = (
                 await self.database.execute_discovery_search(
                     search=search,
                     limit=limit,
@@ -2423,18 +2467,24 @@ async def post_discovery_search(
             )
 
             # Filter results to those that are accessible to the user
-            for data in temp_catalogs_and_collections:
+            for i, (data, hit_token) in enumerate(
+                zip(temp_catalogs_and_collections, hit_tokens)
+            ):
                 # Get access control array for this collection
                 access_control = data["access_control"]
                 data.pop("access_control")
                 # Append collection to list if user has access
                 if int(access_control[-1]) or int(access_control[user_index]):
                     catalogs_and_collections.append(data)
+                    if len(catalogs_and_collections) >= limit:
+                        if i < len(temp_catalogs_and_collections) - 1:
+                            # Extract token from last result
+                            next_token = hit_token
+                            break
 
             # If catalogs_and_collections now less than limit and more results, will need to run search again, giving next_token
             if len(catalogs_and_collections) >= limit or not next_token:
                 # TODO: implement smarter token logic to return token of last returned ES entry
-                next_token = token
                 break
             token = next_token