infrahub-server 1.1.6__py3-none-any.whl → 1.2.0rc0__py3-none-any.whl

infrahub/graphql/queries/resource_manager.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any
 
-from graphene import Field, Float, Int, List, ObjectType, String
+from graphene import Field, Float, Int, List, NonNull, ObjectType, String
 from infrahub_sdk.utils import extract_fields_first_node
 
 from infrahub.core import registry
@@ -21,8 +21,10 @@ from infrahub.pools.number import NumberUtilizationGetter
 if TYPE_CHECKING:
     from graphql import GraphQLResolveInfo
 
+    from infrahub.core.branch import Branch
     from infrahub.core.node import Node
     from infrahub.core.protocols import CoreNode
+    from infrahub.core.timestamp import Timestamp
     from infrahub.database import InfrahubDatabase
     from infrahub.graphql.initialization import GraphqlContext
 
@@ -69,19 +71,21 @@ def _validate_pool_type(pool_id: str, pool: CoreNode | None = None) -> CoreNode:
 
 class PoolAllocated(ObjectType):
     count = Field(Int, required=True, description="The number of allocations within the selected pool.")
-    edges = Field(List(of_type=PoolAllocatedEdge, required=True), required=True)
+    edges = Field(List(of_type=NonNull(PoolAllocatedEdge), required=True), required=True)
 
     @staticmethod
-    async def resolve(  # pylint: disable=unused-argument
-        root: dict,
+    async def resolve(
+        root: dict,  # noqa: ARG004
         info: GraphQLResolveInfo,
         pool_id: str,
         resource_id: str,
         offset: int = 0,
         limit: int = 10,
     ) -> dict:
-        context: GraphqlContext = info.context
-        pool: CoreNode | None = await NodeManager.get_one(id=pool_id, db=context.db, branch=context.branch)
+        graphql_context: GraphqlContext = info.context
+        pool: CoreNode | None = await NodeManager.get_one(
+            id=pool_id, db=graphql_context.db, branch=graphql_context.branch
+        )
 
         fields = await extract_fields_first_node(info=info)
 
@@ -90,14 +94,19 @@ class PoolAllocated(ObjectType):
         match pool.get_kind():
             case InfrahubKind.NUMBERPOOL:
                 return await resolve_number_pool_allocation(
-                    db=context.db, context=context, pool=pool, fields=fields, offset=offset, limit=limit
+                    db=graphql_context.db,
+                    graphql_context=graphql_context,
+                    pool=pool,
+                    fields=fields,
+                    offset=offset,
+                    limit=limit,
                 )
             case InfrahubKind.IPPREFIXPOOL:
                 allocated_kinds.append(InfrahubKind.IPPREFIX)
             case InfrahubKind.IPADDRESSPOOL:
                 allocated_kinds.append(InfrahubKind.IPADDRESS)
 
-        resources = await pool.resources.get_peers(db=context.db)  # type: ignore[attr-defined,union-attr]
+        resources = await pool.resources.get_peers(db=graphql_context.db)  # type: ignore[attr-defined,union-attr]
         if resource_id not in resources:
             raise ValidationError(
                 input_value=f"The selected pool_id={pool_id} doesn't contain the requested resource_id={resource_id}"
@@ -106,8 +115,8 @@ class PoolAllocated(ObjectType):
         resource = resources[resource_id]
 
         query = await IPPrefixUtilization.init(
-            db=context.db,
-            at=context.at,
+            db=graphql_context.db,
+            at=graphql_context.at,
             ip_prefixes=[resource],
             allocated_kinds=allocated_kinds,
             offset=offset,
@@ -115,10 +124,10 @@ class PoolAllocated(ObjectType):
         )
         response: dict[str, Any] = {}
         if "count" in fields:
-            response["count"] = await query.count(db=context.db)
+            response["count"] = await query.count(db=graphql_context.db)
 
         if edges := fields.get("edges"):
-            await query.execute(db=context.db)
+            await query.execute(db=graphql_context.db)
 
             node_fields = edges.get("node", {})
 
@@ -146,9 +155,9 @@ class PoolAllocated(ObjectType):
                 if not identifier_query_class:
                     raise ValidationError(input_value=f"This query doesn't get support {pool.get_kind()}")
                 identifier_query = await identifier_query_class.init(
-                    db=context.db, at=context.at, pool_id=pool_id, allocated=allocated_ids
+                    db=graphql_context.db, at=graphql_context.at, pool_id=pool_id, allocated=allocated_ids
                 )
-                await identifier_query.execute(db=context.db)
+                await identifier_query.execute(db=graphql_context.db)
 
                 reservations = {}
                 for result in identifier_query.get_results():
@@ -171,20 +180,22 @@ class PoolUtilization(ObjectType):
     utilization_default_branch = Field(
         Float, required=True, description="The overall utilization of the pool isolated to the default branch."
     )
-    edges = Field(List(of_type=IPPrefixUtilizationEdge, required=True), required=True)
+    edges = Field(List(of_type=NonNull(IPPrefixUtilizationEdge), required=True), required=True)
 
     @staticmethod
-    async def resolve(  # pylint: disable=unused-argument,too-many-branches
-        root: dict,
+    async def resolve(
+        root: dict,  # noqa: ARG004
         info: GraphQLResolveInfo,
         pool_id: str,
     ) -> dict:
-        context: GraphqlContext = info.context
-        db: InfrahubDatabase = context.db
-        pool: CoreNode | None = await NodeManager.get_one(id=pool_id, db=db, branch=context.branch)
+        graphql_context: GraphqlContext = info.context
+        db: InfrahubDatabase = graphql_context.db
+        pool: CoreNode | None = await NodeManager.get_one(id=pool_id, db=db, branch=graphql_context.branch)
         pool = _validate_pool_type(pool_id=pool_id, pool=pool)
         if pool.get_kind() == "CoreNumberPool":
-            return await resolve_number_pool_utilization(db=db, context=context, pool=pool)
+            return await resolve_number_pool_utilization(
+                db=db, at=graphql_context.at, pool=pool, branch=graphql_context.branch
+            )
 
         resources_map: dict[str, Node] = {}
 
@@ -193,7 +204,9 @@ class PoolUtilization(ObjectType):
         except SchemaNotFoundError:
             pass
 
-        utilization_getter = PrefixUtilizationGetter(db=db, ip_prefixes=list(resources_map.values()), at=context.at)
+        utilization_getter = PrefixUtilizationGetter(
+            db=db, ip_prefixes=list(resources_map.values()), at=graphql_context.at
+        )
         fields = await extract_fields_first_node(info=info)
         response: dict[str, Any] = {}
         total_utilization = None
@@ -262,11 +275,11 @@ class PoolUtilization(ObjectType):
 
 
 async def resolve_number_pool_allocation(
-    db: InfrahubDatabase, context: GraphqlContext, pool: CoreNode, fields: dict, offset: int, limit: int
+    db: InfrahubDatabase, graphql_context: GraphqlContext, pool: CoreNode, fields: dict, offset: int, limit: int
 ) -> dict:
     response: dict[str, Any] = {}
     query = await NumberPoolGetAllocated.init(
-        db=db, pool=pool, offset=offset, limit=limit, branch=context.branch, branch_agnostic=True
+        db=db, pool=pool, offset=offset, limit=limit, branch=graphql_context.branch, branch_agnostic=True
     )
 
     if "count" in fields:
@@ -290,8 +303,10 @@ async def resolve_number_pool_allocation(
     return response
 
 
-async def resolve_number_pool_utilization(db: InfrahubDatabase, context: GraphqlContext, pool: CoreNode) -> dict:
-    number_pool = NumberUtilizationGetter(db=db, pool=pool, at=context.at, branch=context.branch)
+async def resolve_number_pool_utilization(
+    db: InfrahubDatabase, pool: CoreNode, at: Timestamp | str | None, branch: Branch
+) -> dict:
+    number_pool = NumberUtilizationGetter(db=db, pool=pool, at=at, branch=branch)
     await number_pool.load_data()
 
     return {
@@ -322,9 +337,10 @@ InfrahubResourcePoolAllocated = Field(
     limit=Int(required=False),
     offset=Int(required=False),
     resolver=PoolAllocated.resolve,
+    required=True,
 )
 
 
 InfrahubResourcePoolUtilization = Field(
-    PoolUtilization, pool_id=String(required=True), resolver=PoolUtilization.resolve
+    PoolUtilization, pool_id=String(required=True), resolver=PoolUtilization.resolve, required=True
 )

infrahub/graphql/queries/search.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import ipaddress
 from typing import TYPE_CHECKING, Any, Optional
 
-from graphene import Boolean, Field, Int, List, ObjectType, String
+from graphene import Boolean, Field, Int, List, NonNull, ObjectType, String
 from infrahub_sdk.utils import extract_fields_first_node, is_valid_uuid
 
 from infrahub.core.constants import InfrahubKind
@@ -27,7 +27,7 @@ class NodeEdge(ObjectType):
 
 class NodeEdges(ObjectType):
     count = Field(Int, required=True)
-    edges = Field(List(of_type=NodeEdge, required=True), required=False)
+    edges = Field(List(of_type=NonNull(NodeEdge)), required=True)
 
 
 def _collapse_ipv6(s: str) -> str:
@@ -97,13 +97,13 @@ def _collapse_ipv6(s: str) -> str:
 
 
 async def search_resolver(
-    root: dict,  # pylint: disable=unused-argument
+    root: dict,  # noqa: ARG001
     info: GraphQLResolveInfo,
     q: str,
     limit: int = 10,
     partial_match: bool = True,
 ) -> dict[str, Any]:
-    context: GraphqlContext = info.context
+    graphql_context: GraphqlContext = info.context
     response: dict[str, Any] = {}
     results: list[CoreNode] = []
 
@@ -111,7 +111,7 @@ async def search_resolver(
 
     if is_valid_uuid(q):
         matching: Optional[CoreNode] = await NodeManager.get_one(
-            db=context.db, branch=context.branch, at=context.at, id=q
+            db=graphql_context.db, branch=graphql_context.branch, at=graphql_context.at, id=q
         )
         if matching:
             results.append(matching)
@@ -124,8 +124,8 @@ async def search_resolver(
 
         for kind in [InfrahubKind.NODE, InfrahubKind.GENERICGROUP]:
             objs = await NodeManager.query(
-                db=context.db,
-                branch=context.branch,
+                db=graphql_context.db,
+                branch=graphql_context.branch,
                 schema=kind,
                 filters={"any__value": q},
                 limit=limit,
@@ -133,7 +133,7 @@ async def search_resolver(
             )
             results.extend(objs)
 
-    if "edges" in fields and len(results) > 0:
+    if "edges" in fields:
         response["edges"] = [{"node": {"id": obj.id, "kind": obj.get_kind()}} for obj in results]
 
     if "count" in fields:
@@ -148,4 +148,5 @@ InfrahubSearchAnywhere = Field(
     limit=Int(required=False),
     partial_match=Boolean(required=False),
     resolver=search_resolver,
+    required=True,
 )

infrahub/graphql/queries/status.py

@@ -2,11 +2,9 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any
 
-from graphene import Boolean, Field, List, ObjectType, String
+from graphene import Boolean, Field, List, NonNull, ObjectType, String
 from infrahub_sdk.utils import extract_fields_first_node
 
-from infrahub.services import services
-
 if TYPE_CHECKING:
     from graphql import GraphQLResolveInfo
 
@@ -30,7 +28,7 @@ class StatusWorkerEdge(ObjectType):
 
 
 class StatusWorkerEdges(ObjectType):
-    edges = Field(List(of_type=StatusWorkerEdge, required=True), required=True)
+    edges = Field(List(of_type=NonNull(StatusWorkerEdge), required=True), required=True)
 
 
 class Status(ObjectType):
@@ -39,15 +37,18 @@ class Status(ObjectType):
 
 
 async def resolve_status(
-    root: dict,  # pylint: disable=unused-argument
+    root: dict,  # noqa: ARG001
     info: GraphQLResolveInfo,
 ) -> dict:
-    context: GraphqlContext = info.context
-    service = context.service or services.service
+    graphql_context: GraphqlContext = info.context
+    service = graphql_context.service
+    if service is None:
+        raise ValueError("GraphqlContext.service is None")
+
     fields = await extract_fields_first_node(info)
     response: dict[str, Any] = {}
     workers = await service.component.list_workers(
-        branch=str(context.branch.uuid) or context.branch.name, schema_hash=True
+        branch=str(graphql_context.branch.uuid) or graphql_context.branch.name, schema_hash=True
     )
 
     if summary := fields.get("summary"):
@@ -63,4 +64,6 @@ async def resolve_status(
     return response
 
 
-InfrahubStatus = Field(Status, resolver=resolve_status)
+InfrahubStatus = Field(
+    Status, description="Retrieve the status of all infrahub workers.", resolver=resolve_status, required=True
+)

infrahub/graphql/queries/task.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any
 
-from graphene import Field, Int, List, ObjectType, String
+from graphene import Field, Int, List, NonNull, ObjectType, String
 from infrahub_sdk.utils import extract_fields_first_node
 from prefect.client.schemas.objects import StateType
 
@@ -17,12 +17,12 @@ if TYPE_CHECKING:
 
 
 class Tasks(ObjectType):
-    edges = List(TaskNodes)
-    count = Int()
+    edges = List(NonNull(TaskNodes), required=True)
+    count = Int(required=True)
 
     @staticmethod
     async def resolve(
-        root: dict,  # pylint: disable=unused-argument
+        root: dict,  # noqa: ARG004
         info: GraphQLResolveInfo,
         limit: int = 10,
         offset: int = 0,
@@ -49,7 +49,7 @@ class Tasks(ObjectType):
 
     @staticmethod
     async def resolve_branch_status(
-        root: dict,  # pylint: disable=unused-argument
+        root: dict,  # noqa: ARG004
         info: GraphQLResolveInfo,
         branch: str,
     ) -> dict[str, Any]:
@@ -72,11 +72,11 @@ class Tasks(ObjectType):
         limit: int | None = None,
         offset: int | None = None,
     ) -> dict[str, Any]:
-        context: GraphqlContext = info.context
+        graphql_context: GraphqlContext = info.context
         fields = await extract_fields_first_node(info)
 
         prefect_tasks = await PrefectTask.query(
-            db=context.db,
+            db=graphql_context.db,
             fields=fields,
             q=q,
             ids=ids,
@@ -97,7 +97,6 @@ class Tasks(ObjectType):
 
 Task = Field(
     Tasks,
-    resolver=Tasks.resolve,
     limit=Int(required=False),
     offset=Int(required=False),
     related_node__ids=List(String),
@@ -106,11 +105,14 @@ Task = Field(
     workflow=List(String),
     ids=List(String),
     q=String(required=False),
+    resolver=Tasks.resolve,
+    required=True,
 )
 
 TaskBranchStatus = Field(
     Tasks,
-    resolver=Tasks.resolve_branch_status,
     branch=String(required=False),
     description="Return the list of all pending or running tasks that can modify the data, for a given branch",
+    resolver=Tasks.resolve_branch_status,
+    required=True,
 )

infrahub/graphql/resolvers/resolver.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Optional
 
+from graphql.type.definition import GraphQLNonNull
 from infrahub_sdk.utils import extract_fields
 
 from infrahub.core.constants import BranchSupportType, InfrahubKind, RelationshipHierarchyDirection
@@ -17,21 +18,21 @@ from ..types import RELATIONS_PROPERTY_MAP, RELATIONS_PROPERTY_MAP_REVERSED
 if TYPE_CHECKING:
     from graphql import GraphQLResolveInfo
 
-    from infrahub.core.schema import MainSchemaTypes, NodeSchema
+    from infrahub.core.schema import NodeSchema
     from infrahub.graphql.initialization import GraphqlContext
 
 
 async def account_resolver(
-    root: dict,  # pylint: disable=unused-argument
+    root: dict,  # noqa: ARG001
     info: GraphQLResolveInfo,
 ) -> dict:
     fields = await extract_fields(info.field_nodes[0].selection_set)
-    context: GraphqlContext = info.context
+    graphql_context: GraphqlContext = info.context
 
-    async with context.db.start_session() as db:
+    async with graphql_context.db.start_session() as db:
         results = await NodeManager.query(
             schema=InfrahubKind.GENERICACCOUNT,
-            filters={"ids": [context.account_session.account_id]},
+            filters={"ids": [graphql_context.account_session.account_id]},
             fields=fields,
             db=db,
             order=OrderModel(disable=True),
@@ -40,7 +41,9 @@ async def account_resolver(
             account_profile = await results[0].to_graphql(db=db, fields=fields)
             return account_profile
 
-        raise NodeNotFoundError(node_type=InfrahubKind.GENERICACCOUNT, identifier=context.account_session.account_id)
+        raise NodeNotFoundError(
+            node_type=InfrahubKind.GENERICACCOUNT, identifier=graphql_context.account_session.account_id
+        )
 
 
 async def default_resolver(*args: Any, **kwargs) -> dict | list[dict] | None:
@@ -72,14 +75,18 @@ async def default_resolver(*args: Any, **kwargs) -> dict | list[dict] | None:
         raise ValueError(f"expected either 2 or 4 args for default_resolver, got {len(args)}")
 
     # Extract the InfraHub schema by inspecting the GQL Schema
-    node_schema: NodeSchema = info.parent_type.graphene_type._meta.schema
+    node_schema: NodeSchema = (
+        info.parent_type.of_type.graphene_type._meta.schema
+        if isinstance(info.parent_type, GraphQLNonNull)
+        else info.parent_type.graphene_type._meta.schema
+    )
 
     # If the field is an attribute, return its value directly
     if field_name not in node_schema.relationship_names:
         return parent.get(field_name, None)
 
     # Extract the contextual information from the request context
-    context: GraphqlContext = info.context
+    graphql_context: GraphqlContext = info.context
 
     # Extract the name of the fields in the GQL query
     fields = await extract_fields(info.field_nodes[0].selection_set)
@@ -94,7 +101,7 @@ async def default_resolver(*args: Any, **kwargs) -> dict | list[dict] | None:
         if "__" in key and value or key in ["id", "ids"]
     }
 
-    async with context.db.start_session() as db:
+    async with graphql_context.db.start_session() as db:
         objs = await NodeManager.query_peers(
             db=db,
             ids=[parent["id"]],
@@ -102,22 +109,23 @@ async def default_resolver(*args: Any, **kwargs) -> dict | list[dict] | None:
             schema=node_rel,
             filters=filters,
             fields=fields,
-            at=context.at,
-            branch=context.branch,
+            at=graphql_context.at,
+            branch=graphql_context.branch,
             branch_agnostic=node_rel.branch is BranchSupportType.AGNOSTIC,
             fetch_peers=True,
         )
 
         if node_rel.cardinality == "many":
             return [
-                await obj.to_graphql(db=db, fields=fields, related_node_ids=context.related_node_ids) for obj in objs
+                await obj.to_graphql(db=db, fields=fields, related_node_ids=graphql_context.related_node_ids)
+                for obj in objs
             ]
 
         # If cardinality is one
         if not objs:
             return None
 
-        return await objs[0].to_graphql(db=db, fields=fields, related_node_ids=context.related_node_ids)
+        return await objs[0].to_graphql(db=db, fields=fields, related_node_ids=graphql_context.related_node_ids)
 
 
 async def parent_field_name_resolver(parent: dict[str, dict], info: GraphQLResolveInfo) -> dict:
@@ -130,7 +138,7 @@ async def parent_field_name_resolver(parent: dict[str, dict], info: GraphQLResol
 
 
 async def default_paginated_list_resolver(
-    root: dict,  # pylint: disable=unused-argument
+    root: dict,  # noqa: ARG001
     info: GraphQLResolveInfo,
     offset: int | None = None,
     limit: int | None = None,
@@ -138,11 +146,16 @@ async def default_paginated_list_resolver(
     partial_match: bool = False,
     **kwargs: dict[str, Any],
 ) -> dict[str, Any]:
-    schema: MainSchemaTypes = info.return_type.graphene_type._meta.schema
+    schema: NodeSchema = (
+        info.return_type.of_type.graphene_type._meta.schema
+        if isinstance(info.return_type, GraphQLNonNull)
+        else info.return_type.graphene_type._meta.schema
+    )
+
     fields = await extract_selection(info.field_nodes[0], schema=schema)
 
-    context: GraphqlContext = info.context
-    async with context.db.start_session() as db:
+    graphql_context: GraphqlContext = info.context
+    async with graphql_context.db.start_session() as db:
         response: dict[str, Any] = {"edges": []}
         filters = {
             key: value for key, value in kwargs.items() if ("__" in key and value is not None) or key in ("ids", "hfid")
@@ -152,7 +165,11 @@ async def default_paginated_list_resolver(
         node_fields = edges.get("node", {})
 
         permission_set: Optional[dict[str, Any]] = None
-        permissions = await get_permissions(schema=schema, context=context) if context.permissions else None
+        permissions = (
+            await get_permissions(schema=schema, graphql_context=graphql_context)
+            if graphql_context.permissions
+            else None
+        )
         if fields.get("permissions"):
             response["permissions"] = permissions
 
@@ -168,11 +185,11 @@ async def default_paginated_list_resolver(
                 schema=schema,
                 filters=filters or None,
                 fields=node_fields,
-                at=context.at,
-                branch=context.branch,
+                at=graphql_context.at,
+                branch=graphql_context.branch,
                 limit=limit,
                 offset=offset,
-                account=context.account_session,
+                account=graphql_context.account_session,
                 include_source=True,
                 include_owner=True,
                 partial_match=partial_match,
@@ -187,8 +204,8 @@ async def default_paginated_list_resolver(
                     db=db,
                     schema=schema,
                     filters=filters,
-                    at=context.at,
-                    branch=context.branch,
+                    at=graphql_context.at,
+                    branch=graphql_context.branch,
                     partial_match=partial_match,
                 )
 
@@ -199,7 +216,7 @@ async def default_paginated_list_resolver(
                     "node": await obj.to_graphql(
                         db=db,
                         fields=node_fields,
-                        related_node_ids=context.related_node_ids,
+                        related_node_ids=graphql_context.related_node_ids,
                         permissions=permission_set,
                     )
                 }
@@ -211,8 +228,8 @@ async def default_paginated_list_resolver(
 
 
 async def single_relationship_resolver(parent: dict, info: GraphQLResolveInfo, **kwargs: Any) -> dict[str, Any]:
-    context: GraphqlContext = info.context
-    resolver = context.single_relationship_resolver
+    graphql_context: GraphqlContext = info.context
+    resolver = graphql_context.single_relationship_resolver
     return await resolver.resolve(parent=parent, info=info, **kwargs)
 
 
@@ -225,9 +242,13 @@ async def many_relationship_resolver(
     fields by only reusing information below the 'node' key.
     """
     # Extract the InfraHub schema by inspecting the GQL Schema
-    node_schema: NodeSchema = info.parent_type.graphene_type._meta.schema
+    node_schema: NodeSchema = (
+        info.parent_type.of_type.graphene_type._meta.schema
+        if isinstance(info.parent_type, GraphQLNonNull)
+        else info.parent_type.graphene_type._meta.schema
+    )
 
-    context: GraphqlContext = info.context
+    graphql_context: GraphqlContext = info.context
 
     # Extract the name of the fields in the GQL query
     fields = await extract_fields(info.field_nodes[0].selection_set)
@@ -255,7 +276,7 @@ async def many_relationship_resolver(
 
     source_kind = node_schema.kind
 
-    async with context.db.start_session() as db:
+    async with graphql_context.db.start_session() as db:
         ids = [parent["id"]]
         if include_descendants:
             query = await NodeGetHierarchyQuery.init(
@@ -263,8 +284,8 @@ async def many_relationship_resolver(
                 direction=RelationshipHierarchyDirection.DESCENDANTS,
                 node_id=parent["id"],
                 node_schema=node_schema,
-                at=context.at,
-                branch=context.branch,
+                at=graphql_context.at,
+                branch=graphql_context.branch,
             )
             if node_schema.hierarchy:
                 source_kind = node_schema.hierarchy
@@ -279,8 +300,8 @@ async def many_relationship_resolver(
                 source_kind=source_kind,
                 schema=node_rel,
                 filters=filters,
-                at=context.at,
-                branch=context.branch,
+                at=graphql_context.at,
+                branch=graphql_context.branch,
                 branch_agnostic=node_rel.branch is BranchSupportType.AGNOSTIC,
             )
 
@@ -296,8 +317,8 @@ async def many_relationship_resolver(
             fields=node_fields,
             offset=offset,
             limit=limit,
-            at=context.at,
-            branch=context.branch,
+            at=graphql_context.at,
+            branch=graphql_context.branch,
             branch_agnostic=node_rel.branch is BranchSupportType.AGNOSTIC,
             fetch_peers=True,
         )
@@ -305,7 +326,8 @@ async def many_relationship_resolver(
         if not objs:
             return response
         node_graph = [
-            await obj.to_graphql(db=db, fields=node_fields, related_node_ids=context.related_node_ids) for obj in objs
+            await obj.to_graphql(db=db, fields=node_fields, related_node_ids=graphql_context.related_node_ids)
+            for obj in objs
         ]
 
         entries = []
@@ -343,9 +365,13 @@ async def hierarchy_resolver(
     fields by only reusing information below the 'node' key.
     """
     # Extract the InfraHub schema by inspecting the GQL Schema
-    node_schema: NodeSchema = info.parent_type.graphene_type._meta.schema
+    node_schema: NodeSchema = (
+        info.parent_type.of_type.graphene_type._meta.schema
+        if isinstance(info.parent_type, GraphQLNonNull)
+        else info.parent_type.graphene_type._meta.schema
+    )
 
-    context: GraphqlContext = info.context
+    graphql_context: GraphqlContext = info.context
 
     # Extract the name of the fields in the GQL query
     fields = await extract_fields(info.field_nodes[0].selection_set)
@@ -364,7 +390,7 @@ async def hierarchy_resolver(
 
     response: dict[str, Any] = {"edges": [], "count": None}
 
-    async with context.db.start_session() as db:
+    async with graphql_context.db.start_session() as db:
         if "count" in fields:
             response["count"] = await NodeManager.count_hierarchy(
                 db=db,
@@ -372,8 +398,8 @@ async def hierarchy_resolver(
                 direction=direction,
                 node_schema=node_schema,
                 filters=filters,
-                at=context.at,
-                branch=context.branch,
+                at=graphql_context.at,
+                branch=graphql_context.branch,
             )
 
         if not node_fields:
@@ -388,8 +414,8 @@ async def hierarchy_resolver(
             fields=node_fields,
             offset=offset,
             limit=limit,
-            at=context.at,
-            branch=context.branch,
+            at=graphql_context.at,
+            branch=graphql_context.branch,
         )
 
         if not objs: