From 56b8ab995f2e200c482d10c26a5e9a8dd9a082ed Mon Sep 17 00:00:00 2001
From: Onkar Ravgan <onkar.10r@gmail.com>
Date: Fri, 21 Jun 2024 16:34:49 +0530
Subject: [PATCH] MINOR: Add method to list custom properties for a entity for
 python sdk (#16753)

* List custom properties for a entity

* added test

* fixed test
---
 .../ometa/mixins/custom_property_mixin.py     | 15 +++-
 .../ometa/test_ometa_custom_properties_api.py | 75 +++++++++++++++++++
 2 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/ingestion/src/metadata/ingestion/ometa/mixins/custom_property_mixin.py b/ingestion/src/metadata/ingestion/ometa/mixins/custom_property_mixin.py
index 5b0fe01115c0..a2524115600c 100644
--- a/ingestion/src/metadata/ingestion/ometa/mixins/custom_property_mixin.py
+++ b/ingestion/src/metadata/ingestion/ometa/mixins/custom_property_mixin.py
@@ -13,7 +13,9 @@
 
 To be used by OpenMetadata class
 """
-from typing import Dict
+from typing import Dict, List, Optional, Type, TypeVar
+
+from pydantic import BaseModel
 
 from metadata.generated.schema.type.customProperty import PropertyType
 from metadata.generated.schema.type.entityReference import EntityReference
@@ -28,6 +30,8 @@
 
 logger = ometa_logger()
 
+T = TypeVar("T", bound=BaseModel)
+
 
 class OMetaCustomPropertyMixin:
     """
@@ -78,3 +82,12 @@ def get_property_type_ref(self, data_type: CustomPropertyDataTypes) -> PropertyT
         return PropertyType(
             __root__=EntityReference(id=custom_property_type.id, type="type")
         )
+    
+    def get_entity_custom_properties(self, entity_type: Type[T]) -> Optional[List]:
+        """
+        Get all the custom properties of an entity
+        """
+        resp = self.client.get(
+            f"/metadata/types/name/{ENTITY_REFERENCE_TYPE_MAP.get(entity_type.__name__)}?fields=customProperties"
+        )
+        return resp.get("customProperties")
diff --git a/ingestion/tests/integration/ometa/test_ometa_custom_properties_api.py b/ingestion/tests/integration/ometa/test_ometa_custom_properties_api.py
index d2a1854714a1..ec90749f6148 100644
--- a/ingestion/tests/integration/ometa/test_ometa_custom_properties_api.py
+++ b/ingestion/tests/integration/ometa/test_ometa_custom_properties_api.py
@@ -61,6 +61,46 @@
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils.constants import ENTITY_REFERENCE_TYPE_MAP
 
+EXPECTED_CUSTOM_PROPERTIES = [
+    {
+        "name": "DataEngineers",
+        "description": "Data Engineers of a table",
+        "propertyType": {
+            "name": "entityReferenceList",
+        },
+        "customPropertyConfig": {"config": ["user"]},
+    },
+    {
+        "name": "DataQuality",
+        "description": "Quality Details of a Table",
+        "propertyType": {
+            "name": "markdown",
+        },
+    },
+    {
+        "name": "Department",
+        "description": "Department of a table",
+        "propertyType": {
+            "type": "type",
+            "name": "enum",
+        },
+        "customPropertyConfig": {
+            "config": {"values": ["D1", "D2", "D3"], "multiSelect": True}
+        },
+    },
+    {
+        "name": "Rating",
+        "description": "Rating of a table",
+        "propertyType": {"name": "enum"},
+        "customPropertyConfig": {"config": {"values": ["Good", "Average", "Bad"]}},
+    },
+    {
+        "name": "TableSize",
+        "description": "Size of the Table",
+        "propertyType": {"name": "string"},
+    },
+]
+
 
 class OMetaCustomAttributeTest(TestCase):
     """
@@ -277,6 +317,41 @@ def create_custom_property(self):
             ometa_custom_property=ometa_custom_property_request
         )
 
+    def test_get_custom_property(self):
+        """
+        Test getting the custom properties for an entity
+        """
+        # create the custom properties
+        self.create_custom_property()
+
+        custom_properties = self.metadata.get_entity_custom_properties(
+            entity_type=Table
+        )
+
+        actual_custom_properties = []
+        for expected_custom_property in EXPECTED_CUSTOM_PROPERTIES:
+            for custom_property in custom_properties:
+                if expected_custom_property["name"] == custom_property["name"]:
+                    actual_custom_properties.append(custom_property)
+                    self.assertEquals(
+                        custom_property["name"], expected_custom_property["name"]
+                    )
+                    self.assertEquals(
+                        custom_property["description"],
+                        expected_custom_property["description"],
+                    )
+                    self.assertEquals(
+                        custom_property.get("customPropertyConfig"),
+                        expected_custom_property.get("customPropertyConfig"),
+                    )
+                    self.assertEquals(
+                        custom_property["propertyType"]["name"],
+                        expected_custom_property["propertyType"]["name"],
+                    )
+        self.assertEquals(
+            len(actual_custom_properties), len(EXPECTED_CUSTOM_PROPERTIES)
+        )
+
     def test_add_custom_property_table(self):
         """
         Test to add the extension/custom property to the table