From e6915195581021d20c7243e9d7c6578f2f58fe25 Mon Sep 17 00:00:00 2001
From: Murdo <109604278+murdo-moj@users.noreply.github.com>
Date: Wed, 8 Jan 2025 13:34:34 +0000
Subject: [PATCH] Enhance GraphQL queries to includeSoftDelete parameter for
 relationships (#1210)

* Enhance GraphQL queries to includeSoftDelete parameter for relationships

* Enhance getContainer GraphQL query to includeSoftDelete parameter in relationships

* Refactor EntityParser to simplify description retrieval from entity properties

* Refactor EntityParser to streamline description retrieval from entity properties

* Fix EntityParser to handle None properties correctly in description retrieval
---
 .../client/graphql/getChartDetails.graphql         |  8 +++++++-
 .../client/graphql/getContainerDetails.graphql     | 14 ++++++++++++--
 .../client/graphql/getDashboardDetails.graphql     |  8 +++++++-
 .../client/graphql/getDatasetDetails.graphql       |  7 ++++++-
 .../data_platform_catalogue/client/parsers.py      | 12 +++++-------
 5 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/lib/datahub-client/data_platform_catalogue/client/graphql/getChartDetails.graphql b/lib/datahub-client/data_platform_catalogue/client/graphql/getChartDetails.graphql
index c030e55b..65cef020 100644
--- a/lib/datahub-client/data_platform_catalogue/client/graphql/getChartDetails.graphql
+++ b/lib/datahub-client/data_platform_catalogue/client/graphql/getChartDetails.graphql
@@ -6,7 +6,13 @@ query getChartDetails($urn: String!) {
       name
     }
     relationships(
-      input: { types: ["Contains"], direction: INCOMING, start: 0, count: 10 }
+      input: {
+        types: ["Contains"]
+        direction: INCOMING
+        start: 0
+        count: 10
+        includeSoftDelete: false
+      }
     ) {
       total
       relationships {
diff --git a/lib/datahub-client/data_platform_catalogue/client/graphql/getContainerDetails.graphql b/lib/datahub-client/data_platform_catalogue/client/graphql/getContainerDetails.graphql
index fb6f81fe..ef4ae8b5 100644
--- a/lib/datahub-client/data_platform_catalogue/client/graphql/getContainerDetails.graphql
+++ b/lib/datahub-client/data_platform_catalogue/client/graphql/getContainerDetails.graphql
@@ -31,7 +31,15 @@ query getContainer($urn: String!) {
     subTypes {
       typeNames
     }
-    relationships(input: {types: ["IsPartOf"], direction:INCOMING start: 0, count: 500 }) {
+    relationships(
+      input: {
+        types: ["IsPartOf"]
+        direction: INCOMING
+        start: 0
+        count: 500
+        includeSoftDelete: false
+      }
+    ) {
       total
       relationships {
         entity {
@@ -44,7 +52,9 @@ query getContainer($urn: String!) {
             }
             tags {
               tags {
-                tag{urn}
+                tag {
+                  urn
+                }
               }
             }
             subTypes {
diff --git a/lib/datahub-client/data_platform_catalogue/client/graphql/getDashboardDetails.graphql b/lib/datahub-client/data_platform_catalogue/client/graphql/getDashboardDetails.graphql
index 98c81129..2980a4f4 100644
--- a/lib/datahub-client/data_platform_catalogue/client/graphql/getDashboardDetails.graphql
+++ b/lib/datahub-client/data_platform_catalogue/client/graphql/getDashboardDetails.graphql
@@ -23,7 +23,13 @@ query getDashboard($urn: String!) {
       description
     }
     relationships(
-      input: { types: ["Contains"], direction: OUTGOING, start: 0, count: 500 }
+      input: {
+        types: ["Contains"]
+        direction: OUTGOING
+        start: 0
+        count: 500
+        includeSoftDelete: false
+      }
     ) {
       total
       relationships {
diff --git a/lib/datahub-client/data_platform_catalogue/client/graphql/getDatasetDetails.graphql b/lib/datahub-client/data_platform_catalogue/client/graphql/getDatasetDetails.graphql
index 0024ceeb..c2255345 100644
--- a/lib/datahub-client/data_platform_catalogue/client/graphql/getDatasetDetails.graphql
+++ b/lib/datahub-client/data_platform_catalogue/client/graphql/getDatasetDetails.graphql
@@ -57,7 +57,12 @@ query getDatasetDetails($urn: String!) {
       typeNames
     }
     parent_container_relations: relationships(
-      input: { types: ["IsPartOf"], direction: OUTGOING, count: 10 }
+      input: {
+        types: ["IsPartOf"]
+        direction: OUTGOING
+        count: 10
+        includeSoftDelete: false
+      }
     ) {
       total
       relationships {
diff --git a/lib/datahub-client/data_platform_catalogue/client/parsers.py b/lib/datahub-client/data_platform_catalogue/client/parsers.py
index eca455f0..8edd804f 100644
--- a/lib/datahub-client/data_platform_catalogue/client/parsers.py
+++ b/lib/datahub-client/data_platform_catalogue/client/parsers.py
@@ -47,8 +47,6 @@
 
 
 class EntityParser:
-    def __init__(self):
-        pass
 
     def parse(self, search_response) -> SearchResult:
         """Parse graphql response to a SearchResult object"""
@@ -430,11 +428,11 @@ def parse_relations(
                     if relation.get("entity", {}).get("properties") is not None
                     else relation.get("entity").get("name", "")
                 )
-                description = (
-                    relation.get("entity").get("properties", {}).get("description", "")
-                    if relation.get("entity", {}).get("properties") is not None
-                    else ""
-                )
+                properties = relation.get("entity", {}).get("properties")
+                if properties is not None:
+                    description = properties.get("description") or ""
+                elif properties is None:
+                    description = ""
                 tags = self.parse_tags(relation.get("entity"))
                 related_entities.append(
                     EntitySummary(