feat(ingest): add entity_supports_aspect helper (#9120)

metadata-ingestion/src/datahub/emitter/mcp_builder.py

@@ -1,9 +1,10 @@
-from typing import Dict, Iterable, List, Optional, TypeVar
+from typing import Dict, Iterable, List, Optional, Type, TypeVar
 
 from pydantic.fields import Field
 from pydantic.main import BaseModel
 
 from datahub.emitter.mce_builder import (
+    Aspect,
     datahub_guid,
     make_container_urn,
     make_data_platform_urn,
@@ -18,6 +19,7 @@
 )
 from datahub.metadata.com.linkedin.pegasus2avro.container import ContainerProperties
 from datahub.metadata.schema_classes import (
+    KEY_ASPECTS,
     ContainerClass,
     DomainsClass,
     EmbedClass,
@@ -306,3 +308,12 @@ def create_embed_mcp(urn: str, embed_url: str) -> MetadataChangeProposalWrapper:
         entityUrn=urn,
         aspect=EmbedClass(renderUrl=embed_url),
     )
+
+
+def entity_supports_aspect(entity_type: str, aspect_type: Type[Aspect]) -> bool:
+    entity_key_aspect = KEY_ASPECTS[entity_type]
+    aspect_name = aspect_type.get_aspect_name()
+
+    supported_aspects = entity_key_aspect.ASPECT_INFO["entityAspects"]
+
+    return aspect_name in supported_aspects

metadata-ingestion/tests/unit/test_mcp_builder.py

@@ -1,4 +1,5 @@
 import datahub.emitter.mcp_builder as builder
+from datahub.metadata.schema_classes import StatusClass, TelemetryClientIdClass
 
 
 def test_guid_generator():
@@ -83,3 +84,11 @@ def test_guid_generators():
 
     guid = key.guid()
     assert guid == guid_datahub
+
+
+def test_entity_supports_aspect():
+    assert builder.entity_supports_aspect("dataset", StatusClass)
+    assert not builder.entity_supports_aspect("telemetry", StatusClass)
+
+    assert not builder.entity_supports_aspect("dataset", TelemetryClientIdClass)
+    assert builder.entity_supports_aspect("telemetry", TelemetryClientIdClass)