Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
134 changes: 38 additions & 96 deletions infrahub_sdk/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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,
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down
128 changes: 128 additions & 0 deletions infrahub_sdk/diff.py
Original file line number Diff line number Diff line change
@@ -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
Loading