remove nested include/exclude

Author: LucasG0

infrahub_sdk/client.py

@@ -539,6 +539,8 @@ async def _process_nodes_and_relationships(
         response: dict[str, Any],
         schema_kind: str,
         branch: str,
+        prefetch_relationships: bool,
+        include: list[str] | None,
         timeout: int | None = None,
     ) -> ProcessRelationsNode:
         """Processes InfrahubNode and their Relationships from the GraphQL query response.
@@ -547,6 +549,7 @@ async def _process_nodes_and_relationships(
             response (dict[str, Any]): The response from the GraphQL query.
             schema_kind (str): The kind of schema being queried.
             branch (str): The branch name.
+            prefetch_relationships (bool): Flag to indicate whether to prefetch relationship data.
             timeout (int, optional): Overrides default timeout used when querying the graphql API. Specified in seconds.
 
         Returns:
@@ -562,9 +565,14 @@ async def _process_nodes_and_relationships(
             node = await InfrahubNode.from_graphql(client=self, branch=branch, data=item, timeout=timeout)
             nodes.append(node)
 
-            await node._process_relationships(
-                node_data=item, branch=branch, related_nodes=related_nodes, timeout=timeout
-            )
+            if prefetch_relationships or include is not None:
+                await node._process_relationships(
+                    node_data=item,
+                    branch=branch,
+                    related_nodes=related_nodes,
+                    timeout=timeout,
+                    include=include,
+                )
 
         return ProcessRelationsNode(nodes=nodes, related_nodes=related_nodes)
 
@@ -811,7 +819,9 @@ async def process_page(page_offset: int, page_number: int) -> tuple[dict, Proces
                 response=response,
                 schema_kind=schema.kind,
                 branch=branch,
+                prefetch_relationships=prefetch_relationships,
                 timeout=timeout,
+                include=include,
             )
             return response, process_result
 

infrahub_sdk/node/node.py

@@ -730,48 +730,43 @@ async def generate_query_data_node(
                 continue
 
             rel_schema = self._schema.get_relationship(name=rel_name)
+
             if not rel_schema or (not inherited and rel_schema.inherited):
                 continue
 
-            # Don't fetch many attribute/parent relationships unless they are specified in `include`
-            # TODO Why wouldn't we we fetch them if prefetch_relationships is True?
             if (
                 rel_schema.cardinality == RelationshipCardinality.MANY  # type: ignore[union-attr]
                 and rel_schema.kind not in [RelationshipKind.ATTRIBUTE, RelationshipKind.PARENT]  # type: ignore[union-attr]
                 and not (include and rel_name in include)
             ):
                 continue
 
+            peer_data: dict[str, Any] = {}
             should_fetch_relationship = prefetch_relationships or (include is not None and rel_name in include)
             if rel_schema and should_fetch_relationship:
                 peer_schema = await self._client.schema.get(kind=rel_schema.peer, branch=self._branch)
                 peer_node = InfrahubNode(client=self._client, schema=peer_schema, branch=self._branch)
                 peer_data = await peer_node.generate_query_data_node(
-                    include=include,
-                    exclude=exclude,
                     property=property,
                 )
 
-                # TODO is there a reason why we fetch here even with prefetch_relationships == False?
-                if rel_schema.cardinality == "one":
-                    rel_data = RelatedNode._generate_query_data(peer_data=peer_data, property=property)
-                    # Nodes involved in a hierarchy are required to inherit from a common ancestor node, and graphql
-                    # tries to resolve attributes in this ancestor instead of actual node. To avoid
-                    # invalid queries issues when attribute is missing in the common ancestor, we use a fragment
-                    # to explicit actual node kind we are querying.
-                    if rel_schema.kind == RelationshipKind.HIERARCHY:
-                        data_node = rel_data["node"]
-                        rel_data["node"] = {}
-                        rel_data["node"][f"...on {rel_schema.peer}"] = data_node
-                elif rel_schema.cardinality == "many":
-                    rel_data = RelationshipManager._generate_query_data(peer_data=peer_data, property=property)
-                else:
-                    raise ValueError(f"Unknown relationship cardinality {rel_schema.cardinality}")
+            if rel_schema and rel_schema.cardinality == "one":
+                rel_data = RelatedNode._generate_query_data(peer_data=peer_data, property=property)
+                # Nodes involved in a hierarchy are required to inherit from a common ancestor node, and graphql
+                # tries to resolve attributes in this ancestor instead of actual node. To avoid
+                # invalid queries issues when attribute is missing in the common ancestor, we use a fragment
+                # to explicit actual node kind we are querying.
+                if rel_schema.kind == RelationshipKind.HIERARCHY:
+                    data_node = rel_data["node"]
+                    rel_data["node"] = {}
+                    rel_data["node"][f"...on {rel_schema.peer}"] = data_node
+            elif rel_schema and rel_schema.cardinality == "many":
+                rel_data = RelationshipManager._generate_query_data(peer_data=peer_data, property=property)
 
-                data[rel_name] = rel_data
+            data[rel_name] = rel_data
 
-                if insert_alias:
-                    data[rel_name]["@alias"] = f"__alias__{self._schema.kind}__{rel_name}"
+            if insert_alias:
+                data[rel_name]["@alias"] = f"__alias__{self._schema.kind}__{rel_name}"
 
         return data
 
@@ -890,7 +885,12 @@ async def update(
         await self._process_mutation_result(mutation_name=mutation_name, response=response, timeout=timeout)
 
     async def _process_relationships(
-        self, node_data: dict[str, Any], branch: str, related_nodes: list[InfrahubNode], timeout: int | None = None
+        self,
+        node_data: dict[str, Any],
+        branch: str,
+        related_nodes: list[InfrahubNode],
+        timeout: int | None = None,
+        include: list[str] | None = None,
     ) -> None:
         """Processes the Relationships of a InfrahubNode and add Related Nodes to a list.
 
@@ -901,6 +901,8 @@ async def _process_relationships(
             timeout (int, optional): Overrides default timeout used when querying the graphql API. Specified in seconds.
         """
         for rel_name in self._relationships:
+            if include is not None and rel_name not in include:
+                continue
             rel = getattr(self, rel_name)
             if rel and isinstance(rel, RelatedNode):
                 relation = node_data["node"].get(rel_name, None)

tests/integration/test_node.py

@@ -105,10 +105,11 @@ async def test_node_filters_include(
         await car.save(allow_upsert=True)
         assert car.id is not None
 
+        # Clear store, as when we call `owner.peer`, we actually rely on the peer having being stored in store.
+        client.store._branches = {}
         node_after = await client.get(kind=TESTING_CAR, id=car.id)
 
-        with pytest.raises(ValueError):
-            # match=r"Node must have at least one identifier (ID or HFID) to query it."
+        with pytest.raises(NodeNotFoundError, match=f"Unable to find the node '{person_joe.id}' in the store"):
             _ = node_after.owner.peer
 
         assert len(node_after.tags.peers) == 0