acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/sdk/dataset.py

@@ -26,12 +26,14 @@ from datahub.sdk._shared import (
     HasInstitutionalMemory,
     HasOwnership,
     HasPlatformInstance,
+    HasStructuredProperties,
     HasSubtype,
     HasTags,
     HasTerms,
     LinksInputType,
     OwnersInputType,
     ParentContainerInputType,
+    StructuredPropertyInputType,
     TagInputType,
     TagsInputType,
     TermInputType,
@@ -44,6 +46,10 @@ from datahub.sdk.entity import Entity, ExtraAspectsType
 from datahub.utilities.sentinels import Unset, unset
 
 SchemaFieldInputType: TypeAlias = Union[
+    # There is no Enum variant for schema field types because that would force users to do a mapping
+    # to our enum from the raw source type, so additional complexity on their side.
+    # To avoid that, the raw source native type can be provided as a string,
+    # and we will do the mapping internally (in sql_types.py)
     Tuple[str, str],  # (name, type)
     Tuple[str, str, str],  # (name, type, description)
     models.SchemaFieldClass,
@@ -70,6 +76,11 @@ UpstreamLineageInputType: TypeAlias = Union[
     Dict[DatasetUrnOrStr, ColumnLineageMapping],
 ]
 
+ViewDefinitionInputType: TypeAlias = Union[
+    str,
+    models.ViewPropertiesClass,
+]
+
 
 def _parse_upstream_input(
     upstream_input: UpstreamInputType,
@@ -428,6 +439,7 @@ class Dataset(
     HasTags,
     HasTerms,
     HasDomain,
+    HasStructuredProperties,
     Entity,
 ):
     """Represents a dataset in DataHub.
@@ -464,6 +476,7 @@ class Dataset(
         custom_properties: Optional[Dict[str, str]] = None,
         created: Optional[datetime] = None,
         last_modified: Optional[datetime] = None,
+        view_definition: Optional[ViewDefinitionInputType] = None,
         # Standard aspects.
         parent_container: ParentContainerInputType | Unset = unset,
         subtype: Optional[str] = None,
@@ -471,12 +484,12 @@ class Dataset(
         links: Optional[LinksInputType] = None,
         tags: Optional[TagsInputType] = None,
         terms: Optional[TermsInputType] = None,
-        # TODO structured_properties
         domain: Optional[DomainInputType] = None,
-        extra_aspects: ExtraAspectsType = None,
         # Dataset-specific aspects.
         schema: Optional[SchemaFieldsInputType] = None,
         upstreams: Optional[models.UpstreamLineageClass] = None,
+        structured_properties: Optional[StructuredPropertyInputType] = None,
+        extra_aspects: ExtraAspectsType = None,
     ):
         """Initialize a new Dataset instance.
 
@@ -492,6 +505,7 @@ class Dataset(
             custom_properties: Optional dictionary of custom properties.
             created: Optional creation timestamp.
             last_modified: Optional last modification timestamp.
+            view_definition: Optional view definition for the dataset.
             parent_container: Optional parent container for this dataset.
             subtype: Optional subtype of the dataset.
             owners: Optional list of owners.
@@ -533,6 +547,8 @@ class Dataset(
             self.set_created(created)
         if last_modified is not None:
             self.set_last_modified(last_modified)
+        if view_definition is not None:
+            self.set_view_definition(view_definition)
 
         if parent_container is not unset:
             self._set_container(parent_container)
@@ -548,6 +564,9 @@ class Dataset(
             self.set_terms(terms)
         if domain is not None:
             self.set_domain(domain)
+        if structured_properties is not None:
+            for key, value in structured_properties.items():
+                self.set_structured_property(property_urn=key, values=value)
 
     @classmethod
     def _new_from_graph(cls, urn: Urn, current_aspects: models.AspectBag) -> Self:
@@ -711,6 +730,41 @@ class Dataset(
     def set_last_modified(self, last_modified: datetime) -> None:
         self._ensure_dataset_props().lastModified = make_time_stamp(last_modified)
 
+    @property
+    def view_definition(self) -> Optional[models.ViewPropertiesClass]:
+        """Get the view definition of the dataset.
+
+        Under typical usage, this will be present if the subtype is "View".
+
+        Returns:
+            The view definition if set, None otherwise.
+        """
+        return self._get_aspect(models.ViewPropertiesClass)
+
+    def set_view_definition(self, view_definition: ViewDefinitionInputType) -> None:
+        """Set the view definition of the dataset.
+
+        If you're setting a view definition, subtype should typically be set to "view".
+
+        If a string is provided, it will be treated as a SQL view definition. To set
+        a custom language or other properties, provide a ViewPropertiesClass object.
+
+        Args:
+            view_definition: The view definition to set.
+        """
+        if isinstance(view_definition, models.ViewPropertiesClass):
+            self._set_aspect(view_definition)
+        elif isinstance(view_definition, str):
+            self._set_aspect(
+                models.ViewPropertiesClass(
+                    materialized=False,
+                    viewLogic=view_definition,
+                    viewLanguage="SQL",
+                )
+            )
+        else:
+            assert_never(view_definition)
+
     def _schema_dict(self) -> Dict[str, models.SchemaFieldClass]:
         schema_metadata = self._get_aspect(models.SchemaMetadataClass)
         if schema_metadata is None:

datahub/sdk/entity_client.py

@@ -1,15 +1,18 @@
 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,
@@ -19,7 +22,9 @@ 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
@@ -65,6 +70,10 @@ class EntityClient:
     @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.
@@ -84,7 +93,29 @@ 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.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")
@@ -92,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)):
@@ -113,9 +160,12 @@ class EntityClient:
         )
         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.",
@@ -125,9 +175,17 @@ 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)
+        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)
 
@@ -140,7 +198,10 @@ class EntityClient:
         # -> 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)
+        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
@@ -153,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)