From 6f2e6402ea43405dbf2b88204fca525b40252c7c Mon Sep 17 00:00:00 2001
From: Damien Garros <dgarros@gmail.com>
Date: Tue, 31 Dec 2024 12:39:32 +0100
Subject: [PATCH] Fix an issue with session management that could lead to the
 crash of the GraphQL resolver

---
 .../graphql/resolvers/single_relationship.py  |   3 +-
 .../tests/unit/graphql/test_graphql_query.py  | 111 +++++++++++++++++-
 changelog/+gqlquery.fixed.md                  |   1 +
 3 files changed, 113 insertions(+), 2 deletions(-)
 create mode 100644 changelog/+gqlquery.fixed.md

diff --git a/backend/infrahub/graphql/resolvers/single_relationship.py b/backend/infrahub/graphql/resolvers/single_relationship.py
index 5078caf14e..da1a6f4f76 100644
--- a/backend/infrahub/graphql/resolvers/single_relationship.py
+++ b/backend/infrahub/graphql/resolvers/single_relationship.py
@@ -165,4 +165,5 @@ async def _get_entities_with_data_loader(
         node = await loader.load(key=peer_id)
         if not node:
             return None
-        return await node.to_graphql(db=db, fields=node_fields, related_node_ids=related_node_ids)
+        async with db.start_session() as dbs:
+            return await node.to_graphql(db=dbs, fields=node_fields, related_node_ids=related_node_ids)
diff --git a/backend/tests/unit/graphql/test_graphql_query.py b/backend/tests/unit/graphql/test_graphql_query.py
index 578a8f6363..58f904a181 100644
--- a/backend/tests/unit/graphql/test_graphql_query.py
+++ b/backend/tests/unit/graphql/test_graphql_query.py
@@ -9,7 +9,7 @@
 from infrahub.core.constants import BranchSupportType, InfrahubKind
 from infrahub.core.manager import NodeManager
 from infrahub.core.node import Node
-from infrahub.core.schema import NodeSchema
+from infrahub.core.schema import NodeSchema, SchemaRoot
 from infrahub.core.schema.schema_branch import SchemaBranch
 from infrahub.core.timestamp import Timestamp
 from infrahub.database import InfrahubDatabase
@@ -641,6 +641,115 @@ async def test_double_nested_query(db: InfrahubDatabase, default_branch: Branch,
     assert gql_params.context.related_node_ids == {p1.id, p2.id, c1.id, c2.id, c3.id}
 
 
+async def test_nested_query_single_relationship(
+    db: InfrahubDatabase, default_branch: Branch, node_group_schema, data_schema
+):
+    raw_schema = {
+        "version": "1.0",
+        "generics": [
+            {
+                "name": "Generic",
+                "namespace": "Location",
+                "hierarchical": True,
+                "attributes": [{"name": "name", "optional": False, "kind": "Text"}],
+                "relationships": [{"name": "devices", "peer": "InfraDevice", "cardinality": "many", "optional": True}],
+            }
+        ],
+        "nodes": [
+            {
+                "name": "Device",
+                "namespace": "Infra",
+                "attributes": [{"name": "name", "kind": "Text", "optional": False}],
+                "relationships": [
+                    {"name": "location", "peer": "LocationGeneric", "optional": False, "cardinality": "one"}
+                ],
+            },
+            {
+                "name": "Site",
+                "namespace": "Location",
+                "inherit_from": ["LocationGeneric"],
+                "attributes": [{"name": "description", "optional": False, "kind": "Text"}],
+            },
+        ],
+    }
+    schema = SchemaRoot(**raw_schema)
+    schema_branch = registry.schema.register_schema(schema=schema, branch=default_branch.name)
+
+    site_schema = schema_branch.get_node(name="LocationSite")
+    device_schema = schema_branch.get_node(name="InfraDevice")
+
+    site1 = await Node.init(db=db, schema=site_schema, branch=default_branch)
+    await site1.new(db=db, name="site1", description="test")
+    await site1.save(db=db)
+
+    device1 = await Node.init(db=db, schema=device_schema, branch=default_branch)
+    await device1.new(db=db, name="device1", location=site1)
+    await device1.save(db=db)
+
+    device2 = await Node.init(db=db, schema=device_schema, branch=default_branch)
+    await device2.new(db=db, name="device2", location=site1)
+    await device2.save(db=db)
+
+    query = """
+    fragment LocationData on LocationSite {
+        name {
+            value
+        }
+        devices {
+            edges {
+                node {
+                    name {
+                        value
+                    }
+                }
+            }
+        }
+    }
+
+    query {
+        InfraDevice {
+            edges {
+                node {
+                    name {
+                        value
+                    }
+                    location {
+                        node {
+                            ... LocationData
+                        }
+                    }
+                }
+            }
+        }
+    }
+    """
+    gql_params = prepare_graphql_params(
+        db=db, include_mutation=False, include_subscription=False, branch=default_branch
+    )
+    result = await graphql(
+        schema=gql_params.schema,
+        source=query,
+        context_value=gql_params.context,
+        root_value=None,
+        variable_values={},
+    )
+
+    assert result.errors is None
+
+    result_per_name = {
+        result["node"]["name"]["value"]: result["node"] for result in result.data["InfraDevice"]["edges"]
+    }
+    assert sorted(result_per_name.keys()) == ["device1", "device2"]
+    expected_location_data = {
+        "node": {
+            "name": {"value": "site1"},
+            "devices": {"edges": [{"node": {"name": {"value": "device1"}}}, {"node": {"name": {"value": "device2"}}}]},
+        }
+    }
+    assert result.data["InfraDevice"]["edges"][0]["node"]["location"] == expected_location_data
+    assert result.data["InfraDevice"]["edges"][1]["node"]["location"] == expected_location_data
+
+
 async def test_display_label_nested_query(
     db: InfrahubDatabase, default_branch: Branch, car_person_schema: SchemaBranch
 ):
diff --git a/changelog/+gqlquery.fixed.md b/changelog/+gqlquery.fixed.md
new file mode 100644
index 0000000000..0ffdcd3e9f
--- /dev/null
+++ b/changelog/+gqlquery.fixed.md
@@ -0,0 +1 @@
+Fix an issue with session management that could lead to the crash of the GraphQL resolver
\ No newline at end of file