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