acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/sdk/entity_client.py

@@ -1,30 +1,54 @@
 from __future__ import annotations
 
 import warnings
-from typing import TYPE_CHECKING, Union, overload
+from typing import TYPE_CHECKING, Optional, Union, overload
 
 import datahub.metadata.schema_classes as models
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_patch_builder import MetadataPatchProposal
+from datahub.emitter.rest_emitter import EmitMode
 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,
     Urn,
 )
 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
+from datahub.sdk.mlmodelgroup import MLModelGroup
 
 if TYPE_CHECKING:
     from datahub.sdk.main_client import DataHubClient
 
 
 class EntityClient:
+    """Client for managing DataHub entities.
+
+    This class provides methods for retrieving and managing DataHub entities
+    such as datasets, containers, and other metadata objects.
+    """
+
     def __init__(self, client: DataHubClient):
+        """Private constructor - use :py:attr:`DataHubClient.entities` instead.
+
+        Args:
+            client: The parent DataHubClient instance.
+        """
         self._client = client
 
     # TODO: Make all of these methods sync by default.
@@ -38,13 +62,60 @@ class EntityClient:
     @overload
     def get(self, urn: DatasetUrn) -> Dataset: ...
     @overload
+    def get(self, urn: MlModelUrn) -> MLModel: ...
+    @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.
+
+        Args:
+            urn: The urn of the entity to retrieve. Can be a string or :py:class:`Urn` object.
+
+        Returns:
+            The retrieved entity instance.
+
+        Raises:
+            ItemNotFoundError: If the entity does not exist.
+            SdkUsageError: If the entity type is not yet supported.
+            InvalidUrnError: If the URN is invalid.
+        """
         if not isinstance(urn, Urn):
             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.entities.assertion import Assertion
+                from acryl_datahub_cloud.sdk.entities.monitor import Monitor
+                from acryl_datahub_cloud.sdk.entities.subscription import Subscription
+
+                if urn.entity_type == "assertion":
+                    EntityClass = Assertion
+                elif urn.entity_type == "subscription":
+                    EntityClass = Subscription
+                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")
@@ -52,9 +123,25 @@ 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.entities.assertion import Assertion
+
+            assert isinstance(entity, Assertion)
+        elif urn.entity_type == "monitor":
+            from acryl_datahub_cloud.sdk.entities.monitor import Monitor
+
+            assert isinstance(entity, Monitor)
+        elif urn.entity_type == "subscription":
+            from acryl_datahub_cloud.sdk.entities.subscription import Subscription
+
+            assert isinstance(entity, Subscription)
 
-    def create(self, entity: Entity) -> None:
+        return entity
+
+    def create(self, entity: Entity, *, emit_mode: Optional[EmitMode] = None) -> None:
         mcps = []
 
         if self._graph.exists(str(entity.urn)):
@@ -71,11 +158,14 @@ class EntityClient:
                 changeType=models.ChangeTypeClass.CREATE_ENTITY,
             )
         )
-        mcps.extend(entity._as_mcps(models.ChangeTypeClass.CREATE))
+        mcps.extend(entity.as_mcps(models.ChangeTypeClass.CREATE))
 
-        self._graph.emit_mcps(mcps)
+        if emit_mode:
+            self._graph.emit_mcps(mcps, emit_mode=emit_mode)
+        else:
+            self._graph.emit_mcps(mcps)
 
-    def upsert(self, entity: Entity) -> None:
+    def upsert(self, entity: Entity, *, emit_mode: Optional[EmitMode] = None) -> None:
         if entity._prev_aspects is None and self._graph.exists(str(entity.urn)):
             warnings.warn(
                 f"The entity {entity.urn} already exists. This operation will partially overwrite the existing entity.",
@@ -84,10 +174,18 @@ class EntityClient:
             )
             # TODO: If there are no previous aspects but the entity exists, should we delete aspects that are not present here?
 
-        mcps = entity._as_mcps(models.ChangeTypeClass.UPSERT)
-        self._graph.emit_mcps(mcps)
+        mcps = entity.as_mcps(models.ChangeTypeClass.UPSERT)
+        if emit_mode:
+            self._graph.emit_mcps(mcps, emit_mode=emit_mode)
+        else:
+            self._graph.emit_mcps(mcps)
 
-    def update(self, entity: Union[Entity, MetadataPatchProposal]) -> None:
+    def update(
+        self,
+        entity: Union[Entity, MetadataPatchProposal],
+        *,
+        emit_mode: Optional[EmitMode] = None,
+    ) -> None:
         if isinstance(entity, MetadataPatchProposal):
             return self._update_patch(entity)
 
@@ -99,8 +197,11 @@ class EntityClient:
         # TODO: respect If-Unmodified-Since?
         # -> probably add a "mode" parameter that can be "update" (e.g. if not modified) or "update_force"
 
-        mcps = entity._as_mcps(models.ChangeTypeClass.UPSERT)
-        self._graph.emit_mcps(mcps)
+        mcps = entity.as_mcps(models.ChangeTypeClass.UPSERT)
+        if emit_mode:
+            self._graph.emit_mcps(mcps, emit_mode=emit_mode)
+        else:
+            self._graph.emit_mcps(mcps)
 
     def _update_patch(
         self, updater: MetadataPatchProposal, check_exists: bool = True
@@ -113,3 +214,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)