Skip to content

Commit 93b9994

Browse files
ajtmccartyajtmccarty
committed
use DiffSummaryTree instead of deprecated DiffSummary query
1 parent 656b382 commit 93b9994

File tree

2 files changed

+166
-96
lines changed

2 files changed

+166
-96
lines changed

‎infrahub_sdk/client.py‎

Lines changed: 38 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,7 @@
2121

2222
import httpx
2323
import ujson
24-
from typing_extensions import NotRequired, Self
25-
from typing_extensions import TypedDict as ExtensionTypedDict
24+
from typing_extensions import Self
2625

2726
from infrahub_sdk.batch import InfrahubBatch
2827
from infrahub_sdk.branch import (
@@ -33,6 +32,7 @@
3332
from infrahub_sdk.config import Config
3433
from infrahub_sdk.constants import InfrahubClientMode
3534
from infrahub_sdk.data import RepositoryBranchInfo, RepositoryData
35+
from infrahub_sdk.diff import NodeDiff, diff_tree_node_to_node_diff, get_diff_summary_query
3636
from infrahub_sdk.exceptions import (
3737
AuthenticationError,
3838
Error,
@@ -66,34 +66,6 @@
6666
SchemaTypeSync = TypeVar("SchemaTypeSync", bound=CoreNodeSync)
6767

6868

69-
class NodeDiff(ExtensionTypedDict):
70-
branch: str
71-
kind: str
72-
id: str
73-
action: str
74-
display_label: str
75-
elements: list[NodeDiffElement]
76-
77-
78-
class NodeDiffElement(ExtensionTypedDict):
79-
name: str
80-
element_type: str
81-
action: str
82-
summary: NodeDiffSummary
83-
peers: NotRequired[list[NodeDiffPeer]]
84-
85-
86-
class NodeDiffSummary(ExtensionTypedDict):
87-
added: int
88-
updated: int
89-
removed: int
90-
91-
92-
class NodeDiffPeer(ExtensionTypedDict):
93-
action: str
94-
summary: NodeDiffSummary
95-
96-
9769
class ProcessRelationsNode(TypedDict):
9870
nodes: list[InfrahubNode]
9971
related_nodes: list[InfrahubNode]
@@ -1011,41 +983,26 @@ async def get_diff_summary(
1011983
tracker: Optional[str] = None,
1012984
raise_for_error: bool = True,
1013985
) -> list[NodeDiff]:
1014-
query = """
1015-
query {
1016-
DiffSummary {
1017-
branch
1018-
id
1019-
kind
1020-
action
1021-
display_label
1022-
elements {
1023-
element_type
1024-
name
1025-
action
1026-
summary {
1027-
added
1028-
updated
1029-
removed
1030-
}
1031-
... on DiffSummaryElementRelationshipMany {
1032-
peers {
1033-
action
1034-
summary {
1035-
added
1036-
updated
1037-
removed
1038-
}
1039-
}
1040-
}
1041-
}
1042-
}
1043-
}
1044-
"""
986+
query = get_diff_summary_query()
1045987
response = await self.execute_graphql(
1046-
query=query, branch_name=branch, timeout=timeout, tracker=tracker, raise_for_error=raise_for_error
988+
query=query,
989+
branch_name=branch,
990+
timeout=timeout,
991+
tracker=tracker,
992+
raise_for_error=raise_for_error,
993+
variables={"branch_name": branch},
1047994
)
1048-
return response["DiffSummary"]
995+
996+
node_diffs: list[NodeDiff] = []
997+
diff_tree = response["DiffTree"]
998+
999+
if diff_tree is None or "nodes" not in diff_tree:
1000+
return []
1001+
for node_dict in diff_tree["nodes"]:
1002+
node_diff = diff_tree_node_to_node_diff(node_dict=node_dict, branch_name=branch)
1003+
node_diffs.append(node_diff)
1004+
1005+
return node_diffs
10491006

10501007
@overload
10511008
async def allocate_next_ip_address(
@@ -2003,41 +1960,26 @@ def get_diff_summary(
20031960
tracker: Optional[str] = None,
20041961
raise_for_error: bool = True,
20051962
) -> list[NodeDiff]:
2006-
query = """
2007-
query {
2008-
DiffSummary {
2009-
branch
2010-
id
2011-
kind
2012-
action
2013-
display_label
2014-
elements {
2015-
element_type
2016-
name
2017-
action
2018-
summary {
2019-
added
2020-
updated
2021-
removed
2022-
}
2023-
... on DiffSummaryElementRelationshipMany {
2024-
peers {
2025-
action
2026-
summary {
2027-
added
2028-
updated
2029-
removed
2030-
}
2031-
}
2032-
}
2033-
}
2034-
}
2035-
}
2036-
"""
1963+
query = get_diff_summary_query()
20371964
response = self.execute_graphql(
2038-
query=query, branch_name=branch, timeout=timeout, tracker=tracker, raise_for_error=raise_for_error
1965+
query=query,
1966+
branch_name=branch,
1967+
timeout=timeout,
1968+
tracker=tracker,
1969+
raise_for_error=raise_for_error,
1970+
variables={"branch_name": branch},
20391971
)
2040-
return response["DiffSummary"]
1972+
1973+
node_diffs: list[NodeDiff] = []
1974+
diff_tree = response["DiffTree"]
1975+
1976+
if diff_tree is None or "nodes" not in diff_tree:
1977+
return []
1978+
for node_dict in diff_tree["nodes"]:
1979+
node_diff = diff_tree_node_to_node_diff(node_dict=node_dict, branch_name=branch)
1980+
node_diffs.append(node_diff)
1981+
1982+
return node_diffs
20411983

20421984
@overload
20431985
def allocate_next_ip_address(

‎infrahub_sdk/diff.py‎

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
from __future__ import annotations
2+
3+
from typing import (
4+
Any,
5+
)
6+
7+
from typing_extensions import NotRequired, TypedDict
8+
9+
10+
class NodeDiff(TypedDict):
11+
branch: str
12+
kind: str
13+
id: str
14+
action: str
15+
display_label: str
16+
elements: list[NodeDiffElement]
17+
18+
19+
class NodeDiffElement(TypedDict):
20+
name: str
21+
element_type: str
22+
action: str
23+
summary: NodeDiffSummary
24+
peers: NotRequired[list[NodeDiffPeer]]
25+
26+
27+
class NodeDiffSummary(TypedDict):
28+
added: int
29+
updated: int
30+
removed: int
31+
32+
33+
class NodeDiffPeer(TypedDict):
34+
action: str
35+
summary: NodeDiffSummary
36+
37+
38+
def get_diff_summary_query() -> str:
39+
return """
40+
query GetDiffTree($branch_name: String!) {
41+
DiffTree(branch: $branch_name) {
42+
nodes {
43+
uuid
44+
kind
45+
status
46+
label
47+
num_added
48+
num_updated
49+
num_removed
50+
attributes {
51+
name
52+
status
53+
num_added
54+
num_updated
55+
num_removed
56+
}
57+
relationships {
58+
name
59+
status
60+
cardinality
61+
num_added
62+
num_updated
63+
num_removed
64+
elements {
65+
status
66+
num_added
67+
num_updated
68+
num_removed
69+
}
70+
}
71+
}
72+
}
73+
}
74+
"""
75+
76+
77+
def diff_tree_node_to_node_diff(node_dict: dict[str, Any], branch_name: str) -> NodeDiff:
78+
element_diffs: list[NodeDiffElement] = []
79+
if "attributes" in node_dict:
80+
for attr_dict in node_dict["attributes"]:
81+
attr_diff = NodeDiffElement(
82+
action=str(attr_dict.get("status")),
83+
element_type="ATTRIBUTE",
84+
name=str(attr_dict.get("name")),
85+
summary={
86+
"added": int(attr_dict.get("num_added") or 0),
87+
"removed": int(attr_dict.get("num_removed") or 0),
88+
"updated": int(attr_dict.get("num_updated") or 0),
89+
},
90+
)
91+
element_diffs.append(attr_diff)
92+
if "relationships" in node_dict:
93+
for relationship_dict in node_dict["relationships"]:
94+
is_cardinality_one = str(relationship_dict.get("cardinality")).upper() == "ONE"
95+
relationship_diff = NodeDiffElement(
96+
action=str(relationship_dict.get("status")),
97+
element_type="RELATIONSHIP_ONE" if is_cardinality_one else "RELATIONSHIP_MANY",
98+
name=str(relationship_dict.get("name")),
99+
summary={
100+
"added": int(relationship_dict.get("num_added") or 0),
101+
"removed": int(relationship_dict.get("num_removed") or 0),
102+
"updated": int(relationship_dict.get("num_updated") or 0),
103+
},
104+
)
105+
if not is_cardinality_one and "elements" in relationship_dict:
106+
peer_diffs = []
107+
for element_dict in relationship_dict["elements"]:
108+
peer_diffs.append(
109+
NodeDiffPeer(
110+
action=str(element_dict.get("status")),
111+
summary={
112+
"added": int(element_dict.get("num_added") or 0),
113+
"removed": int(element_dict.get("num_removed") or 0),
114+
"updated": int(element_dict.get("num_updated") or 0),
115+
},
116+
)
117+
)
118+
relationship_diff["peers"] = peer_diffs
119+
element_diffs.append(relationship_diff)
120+
node_diff = NodeDiff(
121+
branch=branch_name,
122+
kind=str(node_dict.get("kind")),
123+
id=str(node_dict.get("uuid")),
124+
action=str(node_dict.get("action")),
125+
display_label=str(node_dict.get("label")),
126+
elements=element_diffs,
127+
)
128+
return node_diff

0 commit comments

Comments
 (0)