From d68cf3ac0661a1bfcf99ca75eb905241b05d8b9b Mon Sep 17 00:00:00 2001
From: LavMatt <matthew.laverty@justice.gov.uk>
Date: Wed, 3 Jul 2024 11:46:02 +0100
Subject: [PATCH] return parent as EnityRef in search result

---
 .../data_platform_catalogue/client/search.py  | 23 ++++++++-----------
 1 file changed, 10 insertions(+), 13 deletions(-)

diff --git a/lib/datahub-client/data_platform_catalogue/client/search.py b/lib/datahub-client/data_platform_catalogue/client/search.py
index fc632bb2..3ca37777 100644
--- a/lib/datahub-client/data_platform_catalogue/client/search.py
+++ b/lib/datahub-client/data_platform_catalogue/client/search.py
@@ -17,6 +17,7 @@
     parse_properties,
     parse_tags,
 )
+from data_platform_catalogue.entities import EntityRef
 from data_platform_catalogue.search_types import (
     FacetOption,
     MultiSelectFilter,
@@ -232,24 +233,15 @@ def _parse_result(
         last_modified = parse_last_modified(entity)
         name, display_name, qualified_name = parse_names(entity, properties)
         container = entity.get("container")
+        if container:
+            container_name, container_display_name, container_qualified_name = (
+                parse_names(container, container.get("properties"))
+            )
         domain = parse_domain(entity)
 
         metadata = {
             "owner": owner.display_name,
             "owner_email": owner.email,
-            "parent_container_display_name": (
-                container.get("properties").get("name")
-                if container is not None
-                else None
-            ),
-            "parent_container_fqn": (
-                container.get("properties").get("qualifiedName")
-                if container is not None
-                else None
-            ),
-            "parent_container_urn": (
-                container.get("urn") if container is not None else None
-            ),
             "total_parents": entity.get("relationships", {}).get("total", 0),
             "domain_name": domain.display_name,
             "domain_id": domain.urn,
@@ -270,6 +262,11 @@ def _parse_result(
             name=name,
             display_name=display_name,
             fully_qualified_name=qualified_name,
+            parent_entity=(
+                EntityRef(urn=container.get("urn"), display_name=container_display_name)
+                if container
+                else None
+            ),
             description=properties.get("description", ""),
             metadata=metadata,
             tags=tags,