Merge pull request #411 from opsmill/bkr-fix-count

Add partial_match to the client.count()

Author: BeArchiTek

infrahub_sdk/client.py

@@ -562,18 +562,27 @@ async def count(
         at: Timestamp | None = None,
         branch: str | None = None,
         timeout: int | None = None,
+        partial_match: bool = False,
         **kwargs: Any,
     ) -> int:
         """Return the number of nodes of a given kind."""
-        filters = kwargs
-        schema = await self.schema.get(kind=kind, branch=branch)
+        filters: dict[str, Any] = dict(kwargs)
 
+        if partial_match:
+            filters["partial_match"] = True
+
+        schema = await self.schema.get(kind=kind, branch=branch)
         branch = branch or self.default_branch
         if at:
             at = Timestamp(at)
 
+        data: dict[str, Any] = {
+            "count": None,
+            "@filters": filters,
+        }
+
         response = await self.execute_graphql(
-            query=Query(query={schema.kind: {"count": None, "@filters": filters}}).render(),
+            query=Query(query={schema.kind: data}).render(),
             branch_name=branch,
             at=at,
             timeout=timeout,
@@ -801,7 +810,7 @@ async def process_batch() -> tuple[list[InfrahubNode], list[InfrahubNode]]:
             nodes = []
             related_nodes = []
             batch_process = await self.create_batch()
-            count = await self.count(kind=schema.kind, **filters)
+            count = await self.count(kind=schema.kind, partial_match=partial_match, **filters)
             total_pages = (count + pagination_size - 1) // pagination_size
 
             for page_number in range(1, total_pages + 1):
@@ -1683,18 +1692,27 @@ def count(
         at: Timestamp | None = None,
         branch: str | None = None,
         timeout: int | None = None,
+        partial_match: bool = False,
         **kwargs: Any,
     ) -> int:
         """Return the number of nodes of a given kind."""
-        filters = kwargs
-        schema = self.schema.get(kind=kind, branch=branch)
+        filters: dict[str, Any] = dict(kwargs)
 
+        if partial_match:
+            filters["partial_match"] = True
+
+        schema = self.schema.get(kind=kind, branch=branch)
         branch = branch or self.default_branch
         if at:
             at = Timestamp(at)
 
+        data: dict[str, Any] = {
+            "count": None,
+            "@filters": filters,
+        }
+
         response = self.execute_graphql(
-            query=Query(query={schema.kind: {"count": None, "@filters": filters}}).render(),
+            query=Query(query={schema.kind: data}).render(),
             branch_name=branch,
             at=at,
             timeout=timeout,
@@ -1957,7 +1975,7 @@ def process_batch() -> tuple[list[InfrahubNodeSync], list[InfrahubNodeSync]]:
             related_nodes = []
             batch_process = self.create_batch()
 
-            count = self.count(kind=schema.kind, **filters)
+            count = self.count(kind=schema.kind, partial_match=partial_match, **filters)
             total_pages = (count + pagination_size - 1) // pagination_size
 
             for page_number in range(1, total_pages + 1):