add numMatched, numReturned

Author: jonhealy1

stac_fastapi/core/stac_fastapi/core/core.py

@@ -350,7 +350,7 @@ async def all_collections(
         if datetime:
             parsed_datetime = format_datetime_range(date_str=datetime)
 
-        collections, next_token = await self.database.get_all_collections(
+        collections, next_token, maybe_count = await self.database.get_all_collections(
             token=token,
             limit=limit,
             request=request,
@@ -384,7 +384,12 @@ async def all_collections(
             next_link = PagingLinks(next=next_token, request=request).link_next()
             links.append(next_link)
 
-        return stac_types.Collections(collections=filtered_collections, links=links)
+        return stac_types.Collections(
+            collections=filtered_collections,
+            links=links,
+            numberMatched=maybe_count,
+            numberReturned=len(filtered_collections),
+        )
 
     async def get_collection(
         self, collection_id: str, **kwargs

stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/database_logic.py

@@ -179,7 +179,7 @@ async def get_all_collections(
         filter: Optional[Dict[str, Any]] = None,
         query: Optional[Dict[str, Dict[str, Any]]] = None,
         datetime: Optional[str] = None,
-    ) -> Tuple[List[Dict[str, Any]], Optional[str]]:
+    ) -> Tuple[List[Dict[str, Any]], Optional[str], Optional[int]]:
         """Retrieve a list of collections from Elasticsearch, supporting pagination.
 
         Args:
@@ -317,12 +317,30 @@ async def get_all_collections(
                 else {"bool": {"must": query_parts}}
             )
 
-        # Execute the search
-        response = await self.client.search(
-            index=COLLECTIONS_INDEX,
-            body=body,
+        # Create a copy of the body for count query (without pagination and sorting)
+        count_body = body.copy()
+        if "search_after" in count_body:
+            del count_body["search_after"]
+        count_body["size"] = 0
+
+        # Create async tasks for both search and count
+        search_task = asyncio.create_task(
+            self.client.search(
+                index=COLLECTIONS_INDEX,
+                body=body,
+            )
+        )
+
+        count_task = asyncio.create_task(
+            self.client.count(
+                index=COLLECTIONS_INDEX,
+                body={"query": body.get("query", {"match_all": {}})},
+            )
         )
 
+        # Wait for search task to complete
+        response = await search_task
+
         hits = response["hits"]["hits"]
         collections = [
             self.collection_serializer.db_to_stac(
@@ -337,7 +355,22 @@ async def get_all_collections(
             if next_token_values:
                 next_token = next_token_values[0]
 
-        return collections, next_token
+        # Get the total count of collections
+        matched = (
+            response["hits"]["total"]["value"]
+            if response["hits"]["total"]["relation"] == "eq"
+            else None
+        )
+
+        # If count task is done, use its result
+        if count_task.done():
+            try:
+                matched = count_task.result().get("count")
+            except Exception as e:
+                logger = logging.getLogger(__name__)
+                logger.error(f"Count task failed: {e}")
+
+        return collections, next_token, matched
 
     @staticmethod
     def _apply_collection_datetime_filter(

stac_fastapi/opensearch/stac_fastapi/opensearch/database_logic.py

@@ -163,8 +163,9 @@ async def get_all_collections(
         filter: Optional[Dict[str, Any]] = None,
         query: Optional[Dict[str, Dict[str, Any]]] = None,
         datetime: Optional[str] = None,
-    ) -> Tuple[List[Dict[str, Any]], Optional[str]]:
-        """Retrieve a list of collections from OpenSearch, supporting pagination.
+    ) -> Tuple[List[Dict[str, Any]], Optional[str], Optional[int]]:
+        """
+        Retrieve a list of collections from Opensearch, supporting pagination.
 
         Args:
             token (Optional[str]): The pagination token.
@@ -299,12 +300,30 @@ async def get_all_collections(
                 else {"bool": {"must": query_parts}}
             )
 
-        # Execute the search
-        response = await self.client.search(
-            index=COLLECTIONS_INDEX,
-            body=body,
+        # Create a copy of the body for count query (without pagination and sorting)
+        count_body = body.copy()
+        if "search_after" in count_body:
+            del count_body["search_after"]
+        count_body["size"] = 0
+
+        # Create async tasks for both search and count
+        search_task = asyncio.create_task(
+            self.client.search(
+                index=COLLECTIONS_INDEX,
+                body=body,
+            )
+        )
+
+        count_task = asyncio.create_task(
+            self.client.count(
+                index=COLLECTIONS_INDEX,
+                body={"query": body.get("query", {"match_all": {}})},
+            )
         )
 
+        # Wait for search task to complete
+        response = await search_task
+
         hits = response["hits"]["hits"]
         collections = [
             self.collection_serializer.db_to_stac(
@@ -319,7 +338,22 @@ async def get_all_collections(
             if next_token_values:
                 next_token = next_token_values[0]
 
-        return collections, next_token
+        # Get the total count of collections
+        matched = (
+            response["hits"]["total"]["value"]
+            if response["hits"]["total"]["relation"] == "eq"
+            else None
+        )
+
+        # If count task is done, use its result
+        if count_task.done():
+            try:
+                matched = count_task.result().get("count")
+            except Exception as e:
+                logger = logging.getLogger(__name__)
+                logger.error(f"Count task failed: {e}")
+
+        return collections, next_token, matched
 
     async def get_one_item(self, collection_id: str, item_id: str) -> Dict:
         """Retrieve a single item from the database.

stac_fastapi/tests/api/test_api_search_collections.py

@@ -439,6 +439,7 @@ async def test_collections_query_extension(app_client, txn_client, ctx):
     assert f"{test_prefix}-modis" in found_ids
 
 
+@pytest.mark.asyncio
 async def test_collections_datetime_filter(app_client, load_test_data, txn_client):
     """Test filtering collections by datetime."""
     # Create a test collection with a specific temporal extent
@@ -526,3 +527,74 @@ async def test_collections_datetime_filter(app_client, load_test_data, txn_clien
     found_collections = [c for c in resp_json["collections"] if c["id"] == test_collection_id]
     assert len(found_collections) == 1, f"Expected to find collection {test_collection_id} with open-ended past range to a date within its range"
     """
+
+
+@pytest.mark.asyncio
+async def test_collections_number_matched_returned(app_client, txn_client, ctx):
+    """Verify GET /collections returns correct numberMatched and numberReturned values."""
+    # Create multiple collections with different ids
+    base_collection = ctx.collection
+
+    # Create collections with ids in a specific order to test pagination
+    # Use unique prefixes to avoid conflicts between tests
+    test_prefix = f"count-{uuid.uuid4().hex[:8]}"
+    collection_ids = [f"{test_prefix}-{i}" for i in range(10)]
+
+    for i, coll_id in enumerate(collection_ids):
+        test_collection = base_collection.copy()
+        test_collection["id"] = coll_id
+        test_collection["title"] = f"Test Collection {i}"
+        await create_collection(txn_client, test_collection)
+
+    await refresh_indices(txn_client)
+
+    # Test with limit=5
+    resp = await app_client.get(
+        "/collections",
+        params=[("limit", "5")],
+    )
+    assert resp.status_code == 200
+    resp_json = resp.json()
+
+    # Filter collections to only include the ones we created for this test
+    test_collections = [
+        c for c in resp_json["collections"] if c["id"].startswith(test_prefix)
+    ]
+
+    # Should return 5 collections
+    assert len(test_collections) == 5
+
+    # Check that numberReturned matches the number of collections returned
+    assert resp_json["numberReturned"] == len(resp_json["collections"])
+
+    # Check that numberMatched is greater than or equal to numberReturned
+    # (since there might be other collections in the database)
+    assert resp_json["numberMatched"] >= resp_json["numberReturned"]
+
+    # Check that numberMatched includes at least all our test collections
+    assert resp_json["numberMatched"] >= len(collection_ids)
+
+    # Now test with a query that should match only some collections
+    query = {"id": {"eq": f"{test_prefix}-1"}}
+    resp = await app_client.get(
+        "/collections",
+        params=[("query", json.dumps(query))],
+    )
+    assert resp.status_code == 200
+    resp_json = resp.json()
+
+    # Filter collections to only include the ones we created for this test
+    test_collections = [
+        c for c in resp_json["collections"] if c["id"].startswith(test_prefix)
+    ]
+
+    # Should return only 1 collection
+    assert len(test_collections) == 1
+    assert test_collections[0]["id"] == f"{test_prefix}-1"
+
+    # Check that numberReturned matches the number of collections returned
+    assert resp_json["numberReturned"] == len(resp_json["collections"])
+
+    # Check that numberMatched matches the number of collections that match the query
+    # (should be 1 in this case)
+    assert resp_json["numberMatched"] >= 1