Skip to content

Commit 2254f5f

Browse files
committed
remove nested include/exclude
1 parent 575c6a9 commit 2254f5f

File tree

3 files changed

+41
-26
lines changed

3 files changed

+41
-26
lines changed

infrahub_sdk/client.py

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -539,6 +539,8 @@ async def _process_nodes_and_relationships(
539539
response: dict[str, Any],
540540
schema_kind: str,
541541
branch: str,
542+
prefetch_relationships: bool,
543+
include: list[str] | None,
542544
timeout: int | None = None,
543545
) -> ProcessRelationsNode:
544546
"""Processes InfrahubNode and their Relationships from the GraphQL query response.
@@ -547,6 +549,7 @@ async def _process_nodes_and_relationships(
547549
response (dict[str, Any]): The response from the GraphQL query.
548550
schema_kind (str): The kind of schema being queried.
549551
branch (str): The branch name.
552+
prefetch_relationships (bool): Flag to indicate whether to prefetch relationship data.
550553
timeout (int, optional): Overrides default timeout used when querying the graphql API. Specified in seconds.
551554
552555
Returns:
@@ -562,9 +565,14 @@ async def _process_nodes_and_relationships(
562565
node = await InfrahubNode.from_graphql(client=self, branch=branch, data=item, timeout=timeout)
563566
nodes.append(node)
564567

565-
await node._process_relationships(
566-
node_data=item, branch=branch, related_nodes=related_nodes, timeout=timeout
567-
)
568+
if prefetch_relationships or include is not None:
569+
await node._process_relationships(
570+
node_data=item,
571+
branch=branch,
572+
related_nodes=related_nodes,
573+
timeout=timeout,
574+
include=include,
575+
)
568576

569577
return ProcessRelationsNode(nodes=nodes, related_nodes=related_nodes)
570578

@@ -811,7 +819,9 @@ async def process_page(page_offset: int, page_number: int) -> tuple[dict, Proces
811819
response=response,
812820
schema_kind=schema.kind,
813821
branch=branch,
822+
prefetch_relationships=prefetch_relationships,
814823
timeout=timeout,
824+
include=include,
815825
)
816826
return response, process_result
817827

infrahub_sdk/node/node.py

Lines changed: 25 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -730,18 +730,18 @@ async def generate_query_data_node(
730730
continue
731731

732732
rel_schema = self._schema.get_relationship(name=rel_name)
733+
733734
if not rel_schema or (not inherited and rel_schema.inherited):
734735
continue
735736

736-
# Don't fetch many attribute/parent relationships unless they are specified in `include`
737-
# TODO Why wouldn't we we fetch them if prefetch_relationships is True?
738737
if (
739738
rel_schema.cardinality == RelationshipCardinality.MANY # type: ignore[union-attr]
740739
and rel_schema.kind not in [RelationshipKind.ATTRIBUTE, RelationshipKind.PARENT] # type: ignore[union-attr]
741740
and not (include and rel_name in include)
742741
):
743742
continue
744743

744+
peer_data: dict[str, Any] = {}
745745
should_fetch_relationship = prefetch_relationships or (include is not None and rel_name in include)
746746
if rel_schema and should_fetch_relationship:
747747
peer_schema = await self._client.schema.get(kind=rel_schema.peer, branch=self._branch)
@@ -752,26 +752,23 @@ async def generate_query_data_node(
752752
property=property,
753753
)
754754

755-
# TODO is there a reason why we fetch here even with prefetch_relationships == False?
756-
if rel_schema.cardinality == "one":
757-
rel_data = RelatedNode._generate_query_data(peer_data=peer_data, property=property)
758-
# Nodes involved in a hierarchy are required to inherit from a common ancestor node, and graphql
759-
# tries to resolve attributes in this ancestor instead of actual node. To avoid
760-
# invalid queries issues when attribute is missing in the common ancestor, we use a fragment
761-
# to explicit actual node kind we are querying.
762-
if rel_schema.kind == RelationshipKind.HIERARCHY:
763-
data_node = rel_data["node"]
764-
rel_data["node"] = {}
765-
rel_data["node"][f"...on {rel_schema.peer}"] = data_node
766-
elif rel_schema.cardinality == "many":
767-
rel_data = RelationshipManager._generate_query_data(peer_data=peer_data, property=property)
768-
else:
769-
raise ValueError(f"Unknown relationship cardinality {rel_schema.cardinality}")
755+
if rel_schema and rel_schema.cardinality == "one":
756+
rel_data = RelatedNode._generate_query_data(peer_data=peer_data, property=property)
757+
# Nodes involved in a hierarchy are required to inherit from a common ancestor node, and graphql
758+
# tries to resolve attributes in this ancestor instead of actual node. To avoid
759+
# invalid queries issues when attribute is missing in the common ancestor, we use a fragment
760+
# to explicit actual node kind we are querying.
761+
if rel_schema.kind == RelationshipKind.HIERARCHY:
762+
data_node = rel_data["node"]
763+
rel_data["node"] = {}
764+
rel_data["node"][f"...on {rel_schema.peer}"] = data_node
765+
elif rel_schema and rel_schema.cardinality == "many":
766+
rel_data = RelationshipManager._generate_query_data(peer_data=peer_data, property=property)
770767

771-
data[rel_name] = rel_data
768+
data[rel_name] = rel_data
772769

773-
if insert_alias:
774-
data[rel_name]["@alias"] = f"__alias__{self._schema.kind}__{rel_name}"
770+
if insert_alias:
771+
data[rel_name]["@alias"] = f"__alias__{self._schema.kind}__{rel_name}"
775772

776773
return data
777774

@@ -890,7 +887,12 @@ async def update(
890887
await self._process_mutation_result(mutation_name=mutation_name, response=response, timeout=timeout)
891888

892889
async def _process_relationships(
893-
self, node_data: dict[str, Any], branch: str, related_nodes: list[InfrahubNode], timeout: int | None = None
890+
self,
891+
node_data: dict[str, Any],
892+
branch: str,
893+
related_nodes: list[InfrahubNode],
894+
timeout: int | None = None,
895+
include: list[str] | None = None,
894896
) -> None:
895897
"""Processes the Relationships of a InfrahubNode and add Related Nodes to a list.
896898
@@ -901,6 +903,8 @@ async def _process_relationships(
901903
timeout (int, optional): Overrides default timeout used when querying the graphql API. Specified in seconds.
902904
"""
903905
for rel_name in self._relationships:
906+
if include is not None and rel_name not in include:
907+
continue
904908
rel = getattr(self, rel_name)
905909
if rel and isinstance(rel, RelatedNode):
906910
relation = node_data["node"].get(rel_name, None)

tests/integration/test_node.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -105,10 +105,11 @@ async def test_node_filters_include(
105105
await car.save(allow_upsert=True)
106106
assert car.id is not None
107107

108+
# Clear store, as when we call `owner.peer`, we actually rely on the peer having being stored in store.
109+
client.store._branches = {}
108110
node_after = await client.get(kind=TESTING_CAR, id=car.id)
109111

110-
with pytest.raises(ValueError):
111-
# match=r"Node must have at least one identifier (ID or HFID) to query it."
112+
with pytest.raises(NodeNotFoundError, match=f"Unable to find the node '{person_joe.id}' in the store"):
112113
_ = node_after.owner.peer
113114

114115
assert len(node_after.tags.peers) == 0

0 commit comments

Comments
 (0)