diff --git a/infrahub_sdk/client.py b/infrahub_sdk/client.py index 2b2ab0b3..29310ac3 100644 --- a/infrahub_sdk/client.py +++ b/infrahub_sdk/client.py @@ -21,8 +21,7 @@ import httpx import ujson -from typing_extensions import NotRequired, Self -from typing_extensions import TypedDict as ExtensionTypedDict +from typing_extensions import Self from infrahub_sdk.batch import InfrahubBatch from infrahub_sdk.branch import ( @@ -33,6 +32,7 @@ from infrahub_sdk.config import Config from infrahub_sdk.constants import InfrahubClientMode from infrahub_sdk.data import RepositoryBranchInfo, RepositoryData +from infrahub_sdk.diff import NodeDiff, diff_tree_node_to_node_diff, get_diff_summary_query from infrahub_sdk.exceptions import ( AuthenticationError, Error, @@ -66,34 +66,6 @@ SchemaTypeSync = TypeVar("SchemaTypeSync", bound=CoreNodeSync) -class NodeDiff(ExtensionTypedDict): - branch: str - kind: str - id: str - action: str - display_label: str - elements: list[NodeDiffElement] - - -class NodeDiffElement(ExtensionTypedDict): - name: str - element_type: str - action: str - summary: NodeDiffSummary - peers: NotRequired[list[NodeDiffPeer]] - - -class NodeDiffSummary(ExtensionTypedDict): - added: int - updated: int - removed: int - - -class NodeDiffPeer(ExtensionTypedDict): - action: str - summary: NodeDiffSummary - - class ProcessRelationsNode(TypedDict): nodes: list[InfrahubNode] related_nodes: list[InfrahubNode] @@ -1011,41 +983,26 @@ async def get_diff_summary( tracker: Optional[str] = None, raise_for_error: bool = True, ) -> list[NodeDiff]: - query = """ - query { - DiffSummary { - branch - id - kind - action - display_label - elements { - element_type - name - action - summary { - added - updated - removed - } - ... on DiffSummaryElementRelationshipMany { - peers { - action - summary { - added - updated - removed - } - } - } - } - } - } - """ + query = get_diff_summary_query() response = await self.execute_graphql( - query=query, branch_name=branch, timeout=timeout, tracker=tracker, raise_for_error=raise_for_error + query=query, + branch_name=branch, + timeout=timeout, + tracker=tracker, + raise_for_error=raise_for_error, + variables={"branch_name": branch}, ) - return response["DiffSummary"] + + node_diffs: list[NodeDiff] = [] + diff_tree = response["DiffTree"] + + if diff_tree is None or "nodes" not in diff_tree: + return [] + for node_dict in diff_tree["nodes"]: + node_diff = diff_tree_node_to_node_diff(node_dict=node_dict, branch_name=branch) + node_diffs.append(node_diff) + + return node_diffs @overload async def allocate_next_ip_address( @@ -2003,41 +1960,26 @@ def get_diff_summary( tracker: Optional[str] = None, raise_for_error: bool = True, ) -> list[NodeDiff]: - query = """ - query { - DiffSummary { - branch - id - kind - action - display_label - elements { - element_type - name - action - summary { - added - updated - removed - } - ... on DiffSummaryElementRelationshipMany { - peers { - action - summary { - added - updated - removed - } - } - } - } - } - } - """ + query = get_diff_summary_query() response = self.execute_graphql( - query=query, branch_name=branch, timeout=timeout, tracker=tracker, raise_for_error=raise_for_error + query=query, + branch_name=branch, + timeout=timeout, + tracker=tracker, + raise_for_error=raise_for_error, + variables={"branch_name": branch}, ) - return response["DiffSummary"] + + node_diffs: list[NodeDiff] = [] + diff_tree = response["DiffTree"] + + if diff_tree is None or "nodes" not in diff_tree: + return [] + for node_dict in diff_tree["nodes"]: + node_diff = diff_tree_node_to_node_diff(node_dict=node_dict, branch_name=branch) + node_diffs.append(node_diff) + + return node_diffs @overload def allocate_next_ip_address( diff --git a/infrahub_sdk/diff.py b/infrahub_sdk/diff.py new file mode 100644 index 00000000..c445000f --- /dev/null +++ b/infrahub_sdk/diff.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +from typing import ( + Any, +) + +from typing_extensions import NotRequired, TypedDict + + +class NodeDiff(TypedDict): + branch: str + kind: str + id: str + action: str + display_label: str + elements: list[NodeDiffElement] + + +class NodeDiffElement(TypedDict): + name: str + element_type: str + action: str + summary: NodeDiffSummary + peers: NotRequired[list[NodeDiffPeer]] + + +class NodeDiffSummary(TypedDict): + added: int + updated: int + removed: int + + +class NodeDiffPeer(TypedDict): + action: str + summary: NodeDiffSummary + + +def get_diff_summary_query() -> str: + return """ + query GetDiffTree($branch_name: String!) { + DiffTree(branch: $branch_name) { + nodes { + uuid + kind + status + label + num_added + num_updated + num_removed + attributes { + name + status + num_added + num_updated + num_removed + } + relationships { + name + status + cardinality + num_added + num_updated + num_removed + elements { + status + num_added + num_updated + num_removed + } + } + } + } + } + """ + + +def diff_tree_node_to_node_diff(node_dict: dict[str, Any], branch_name: str) -> NodeDiff: + element_diffs: list[NodeDiffElement] = [] + if "attributes" in node_dict: + for attr_dict in node_dict["attributes"]: + attr_diff = NodeDiffElement( + action=str(attr_dict.get("status")), + element_type="ATTRIBUTE", + name=str(attr_dict.get("name")), + summary={ + "added": int(attr_dict.get("num_added") or 0), + "removed": int(attr_dict.get("num_removed") or 0), + "updated": int(attr_dict.get("num_updated") or 0), + }, + ) + element_diffs.append(attr_diff) + if "relationships" in node_dict: + for relationship_dict in node_dict["relationships"]: + is_cardinality_one = str(relationship_dict.get("cardinality")).upper() == "ONE" + relationship_diff = NodeDiffElement( + action=str(relationship_dict.get("status")), + element_type="RELATIONSHIP_ONE" if is_cardinality_one else "RELATIONSHIP_MANY", + name=str(relationship_dict.get("name")), + summary={ + "added": int(relationship_dict.get("num_added") or 0), + "removed": int(relationship_dict.get("num_removed") or 0), + "updated": int(relationship_dict.get("num_updated") or 0), + }, + ) + if not is_cardinality_one and "elements" in relationship_dict: + peer_diffs = [] + for element_dict in relationship_dict["elements"]: + peer_diffs.append( + NodeDiffPeer( + action=str(element_dict.get("status")), + summary={ + "added": int(element_dict.get("num_added") or 0), + "removed": int(element_dict.get("num_removed") or 0), + "updated": int(element_dict.get("num_updated") or 0), + }, + ) + ) + relationship_diff["peers"] = peer_diffs + element_diffs.append(relationship_diff) + node_diff = NodeDiff( + branch=branch_name, + kind=str(node_dict.get("kind")), + id=str(node_dict.get("uuid")), + action=str(node_dict.get("action")), + display_label=str(node_dict.get("label")), + elements=element_diffs, + ) + return node_diff diff --git a/tests/unit/sdk/test_diff_summary.py b/tests/unit/sdk/test_diff_summary.py new file mode 100644 index 00000000..6ab1da99 --- /dev/null +++ b/tests/unit/sdk/test_diff_summary.py @@ -0,0 +1,160 @@ +import pytest +from pytest_httpx import HTTPXMock + +from infrahub_sdk import InfrahubClient +from tests.unit.sdk.conftest import BothClients + +client_types = ["standard", "sync"] + + +@pytest.fixture +async def mock_diff_tree_query(httpx_mock: HTTPXMock, client: InfrahubClient) -> HTTPXMock: + response = { + "data": { + "DiffTree": { + "nodes": [ + { + "attributes": [], + "kind": "TestCar", + "label": "nolt #444444", + "num_added": 0, + "num_removed": 0, + "num_updated": 1, + "relationships": [ + { + "cardinality": "ONE", + "elements": [{"num_added": 0, "num_removed": 0, "num_updated": 1, "status": "UPDATED"}], + "name": "owner", + "num_added": 0, + "num_removed": 0, + "num_updated": 1, + "status": "UPDATED", + } + ], + "status": "UPDATED", + "uuid": "17fbadf0-6637-4fa2-43e6-1677ea170e0f", + }, + { + "attributes": [], + "kind": "TestPerson", + "label": "Jane", + "num_added": 0, + "num_removed": 0, + "num_updated": 1, + "relationships": [ + { + "cardinality": "MANY", + "elements": [{"num_added": 0, "num_removed": 3, "num_updated": 0, "status": "REMOVED"}], + "name": "cars", + "num_added": 0, + "num_removed": 1, + "num_updated": 0, + "status": "UPDATED", + } + ], + "status": "UPDATED", + "uuid": "17fbadf0-634f-05a8-43e4-1677e744d4c0", + }, + { + "attributes": [ + {"name": "name", "num_added": 0, "num_removed": 0, "num_updated": 1, "status": "UPDATED"} + ], + "kind": "TestPerson", + "label": "Jonathan", + "num_added": 0, + "num_removed": 0, + "num_updated": 2, + "relationships": [ + { + "cardinality": "MANY", + "elements": [{"num_added": 3, "num_removed": 0, "num_updated": 0, "status": "ADDED"}], + "name": "cars", + "num_added": 1, + "num_removed": 0, + "num_updated": 0, + "status": "UPDATED", + } + ], + "status": "UPDATED", + "uuid": "17fbadf0-6243-5d3c-43ee-167718ff8dac", + }, + ] + } + } + } + + httpx_mock.add_response( + method="POST", + json=response, + match_headers={"X-Infrahub-Tracker": "query-difftree"}, + ) + return httpx_mock + + +@pytest.mark.parametrize("client_type", client_types) +async def test_diffsummary(clients: BothClients, mock_diff_tree_query, client_type): + if client_type == "standard": + node_diffs = await clients.standard.get_diff_summary( + branch="branch2", + tracker="query-difftree", + ) + else: + node_diffs = clients.sync.get_diff_summary( + branch="branch2", + tracker="query-difftree", + ) + + assert len(node_diffs) == 3 + assert { + "branch": "branch2", + "kind": "TestCar", + "id": "17fbadf0-6637-4fa2-43e6-1677ea170e0f", + "action": "None", + "display_label": "nolt #444444", + "elements": [ + { + "action": "UPDATED", + "element_type": "RELATIONSHIP_ONE", + "name": "owner", + "summary": {"added": 0, "removed": 0, "updated": 1}, + } + ], + } in node_diffs + assert { + "branch": "branch2", + "kind": "TestPerson", + "id": "17fbadf0-634f-05a8-43e4-1677e744d4c0", + "action": "None", + "display_label": "Jane", + "elements": [ + { + "action": "UPDATED", + "element_type": "RELATIONSHIP_MANY", + "name": "cars", + "summary": {"added": 0, "removed": 1, "updated": 0}, + "peers": [{"action": "REMOVED", "summary": {"added": 0, "removed": 3, "updated": 0}}], + } + ], + } in node_diffs + assert { + "branch": "branch2", + "kind": "TestPerson", + "id": "17fbadf0-6243-5d3c-43ee-167718ff8dac", + "action": "None", + "display_label": "Jonathan", + "elements": [ + { + "action": "UPDATED", + "element_type": "ATTRIBUTE", + "name": "name", + "summary": {"added": 0, "removed": 0, "updated": 1}, + }, + { + "action": "UPDATED", + "element_type": "RELATIONSHIP_MANY", + "name": "cars", + "summary": {"added": 1, "removed": 0, "updated": 0}, + "peers": [{"action": "ADDED", "summary": {"added": 3, "removed": 0, "updated": 0}}], + }, + ], + } in node_diffs