acryl-datahub 1.1.0rc4__py3-none-any.whl → 1.1.0.1__py3-none-any.whl

datahub/sdk/entity_client.py

@@ -9,7 +9,11 @@ from datahub.emitter.mcp_patch_builder import MetadataPatchProposal
 from datahub.errors import IngestionAttributionWarning, ItemNotFoundError, SdkUsageError
 from datahub.ingestion.graph.client import DataHubGraph
 from datahub.metadata.urns import (
+    ChartUrn,
     ContainerUrn,
+    DashboardUrn,
+    DataFlowUrn,
+    DataJobUrn,
     DatasetUrn,
     MlModelGroupUrn,
     MlModelUrn,
@@ -17,7 +21,11 @@ from datahub.metadata.urns import (
 )
 from datahub.sdk._all_entities import ENTITY_CLASSES
 from datahub.sdk._shared import UrnOrStr
+from datahub.sdk.chart import Chart
 from datahub.sdk.container import Container
+from datahub.sdk.dashboard import Dashboard
+from datahub.sdk.dataflow import DataFlow
+from datahub.sdk.datajob import DataJob
 from datahub.sdk.dataset import Dataset
 from datahub.sdk.entity import Entity
 from datahub.sdk.mlmodel import MLModel
@@ -57,6 +65,14 @@ class EntityClient:
     @overload
     def get(self, urn: MlModelGroupUrn) -> MLModelGroup: ...
     @overload
+    def get(self, urn: DataFlowUrn) -> DataFlow: ...
+    @overload
+    def get(self, urn: DataJobUrn) -> DataJob: ...
+    @overload
+    def get(self, urn: DashboardUrn) -> Dashboard: ...
+    @overload
+    def get(self, urn: ChartUrn) -> Chart: ...
+    @overload
     def get(self, urn: Union[Urn, str]) -> Entity: ...
     def get(self, urn: UrnOrStr) -> Entity:
         """Retrieve an entity by its urn.
@@ -76,7 +92,26 @@ class EntityClient:
             urn = Urn.from_string(urn)
 
         # TODO: add error handling around this with a suggested alternative if not yet supported
-        EntityClass = ENTITY_CLASSES[urn.entity_type]
+        try:
+            EntityClass = ENTITY_CLASSES[urn.entity_type]
+        except KeyError as e:
+            # Try to import cloud-specific entities if not found
+            try:
+                from acryl_datahub_cloud._sdk_extras.entities.assertion import Assertion
+                from acryl_datahub_cloud._sdk_extras.entities.monitor import Monitor
+
+                if urn.entity_type == "assertion":
+                    EntityClass = Assertion
+                elif urn.entity_type == "monitor":
+                    EntityClass = Monitor
+                else:
+                    raise SdkUsageError(
+                        f"Entity type {urn.entity_type} is not yet supported"
+                    ) from e
+            except ImportError as e:
+                raise SdkUsageError(
+                    f"Entity type {urn.entity_type} is not yet supported"
+                ) from e
 
         if not self._graph.exists(str(urn)):
             raise ItemNotFoundError(f"Entity {urn} not found")
@@ -84,7 +119,19 @@ class EntityClient:
         aspects = self._graph.get_entity_semityped(str(urn))
 
         # TODO: save the timestamp so we can use If-Unmodified-Since on the updates
-        return EntityClass._new_from_graph(urn, aspects)
+        entity = EntityClass._new_from_graph(urn, aspects)
+
+        # Type narrowing for cloud-specific entities
+        if urn.entity_type == "assertion":
+            from acryl_datahub_cloud._sdk_extras.entities.assertion import Assertion
+
+            assert isinstance(entity, Assertion)
+        elif urn.entity_type == "monitor":
+            from acryl_datahub_cloud._sdk_extras.entities.monitor import Monitor
+
+            assert isinstance(entity, Monitor)
+
+        return entity
 
     def create(self, entity: Entity) -> None:
         mcps = []
@@ -145,3 +192,44 @@ class EntityClient:
 
         mcps = updater.build()
         self._graph.emit_mcps(mcps)
+
+    def delete(
+        self,
+        urn: UrnOrStr,
+        check_exists: bool = True,
+        cascade: bool = False,
+        hard: bool = False,
+    ) -> None:
+        """Delete an entity by its urn.
+
+        Args:
+            urn: The urn of the entity to delete. Can be a string or :py:class:`Urn` object.
+            check_exists: Whether to check if the entity exists before deletion. Defaults to True.
+            cascade: Whether to cascade delete related entities. When True, deletes child entities
+                like datajobs within dataflows, datasets within containers, etc. Not yet supported.
+            hard: Whether to perform a hard delete (permanent) or soft delete. Defaults to False.
+
+        Raises:
+            SdkUsageError: If the entity does not exist and check_exists is True, or if cascade is True (not supported).
+
+        Note:
+            When hard is True, the operation is irreversible and the entity will be permanently removed.
+
+            Impact of cascade deletion (still to be done) depends on the input entity type:
+            - Container: Recursively deletes all containers and data assets within the container.
+            - Dataflow: Recursively deletes all data jobs within the dataflow.
+            - Dashboard: TBD
+            - DataPlatformInstance: TBD
+            - ...
+        """
+        urn_str = str(urn) if isinstance(urn, Urn) else urn
+        if check_exists and not self._graph.exists(entity_urn=urn_str):
+            raise SdkUsageError(
+                f"Entity {urn_str} does not exist, and hence cannot be deleted. "
+                "You can bypass this check by setting check_exists=False."
+            )
+
+        if cascade:
+            raise SdkUsageError("The 'cascade' parameter is not yet supported.")
+
+        self._graph.delete_entity(urn=urn_str, hard=hard)