Add count method to clients (#225)

Author: gmazoyer

changelog/158.added.me

@@ -0,0 +1 @@
+Add `count` method to both sync and async clients to retrieve the number of objects of a given kind

infrahub_sdk/client.py

@@ -525,6 +525,25 @@ async def _process_nodes_and_relationships(
 
         return ProcessRelationsNode(nodes=nodes, related_nodes=related_nodes)
 
+    async def count(
+        self,
+        kind: str | type[SchemaType],
+        at: Timestamp | None = None,
+        branch: str | None = None,
+        timeout: int | None = None,
+    ) -> int:
+        """Return the number of nodes of a given kind."""
+        schema = await self.schema.get(kind=kind, branch=branch)
+
+        branch = branch or self.default_branch
+        if at:
+            at = Timestamp(at)
+
+        response = await self.execute_graphql(
+            query=Query(query={schema.kind: {"count": None}}).render(), branch_name=branch, at=at, timeout=timeout
+        )
+        return int(response.get(schema.kind, {}).get("count", 0))
+
     @overload
     async def all(
         self,
@@ -1549,6 +1568,25 @@ def execute_graphql(
 
         # TODO add a special method to execute mutation that will check if the method returned OK
 
+    def count(
+        self,
+        kind: str | type[SchemaType],
+        at: Timestamp | None = None,
+        branch: str | None = None,
+        timeout: int | None = None,
+    ) -> int:
+        """Return the number of nodes of a given kind."""
+        schema = self.schema.get(kind=kind, branch=branch)
+
+        branch = branch or self.default_branch
+        if at:
+            at = Timestamp(at)
+
+        response = self.execute_graphql(
+            query=Query(query={schema.kind: {"count": None}}).render(), branch_name=branch, at=at, timeout=timeout
+        )
+        return int(response.get(schema.kind, {}).get("count", 0))
+
     @overload
     def all(
         self,

tests/unit/sdk/conftest.py

@@ -1680,6 +1680,12 @@ async def mock_query_corenode_page1_1(httpx_mock: HTTPXMock, client: InfrahubCli
     return httpx_mock
 
 
+@pytest.fixture
+async def mock_query_repository_count(httpx_mock: HTTPXMock, client: InfrahubClient, mock_schema_query_01) -> HTTPXMock:
+    httpx_mock.add_response(method="POST", json={"data": {"CoreRepository": {"count": 5}}})
+    return httpx_mock
+
+
 @pytest.fixture
 async def mock_query_repository_page1_empty(
     httpx_mock: HTTPXMock, client: InfrahubClient, mock_schema_query_01

tests/unit/sdk/test_client.py

@@ -66,6 +66,16 @@ async def test_get_repositories(
     }
 
 
+@pytest.mark.parametrize("client_type", client_types)
+async def test_method_count(clients, mock_query_repository_count, client_type):  # pylint: disable=unused-argument
+    if client_type == "standard":
+        count = await clients.standard.count(kind="CoreRepository")
+    else:
+        count = clients.sync.count(kind="CoreRepository")
+
+    assert count == 5
+
+
 @pytest.mark.parametrize("client_type", client_types)
 async def test_method_all_with_limit(clients, mock_query_repository_page1_2, client_type):  # pylint: disable=unused-argument
     if client_type == "standard":