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

datahub/sdk/_shared.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
 import warnings
+from abc import ABC, abstractmethod
 from datetime import datetime
 from typing import (
     TYPE_CHECKING,
     Callable,
+    Dict,
     List,
     Optional,
     Sequence,
@@ -14,6 +16,7 @@ from typing import (
 
 from typing_extensions import TypeAlias, assert_never
 
+import datahub.emitter.mce_builder as builder
 import datahub.metadata.schema_classes as models
 from datahub.emitter.mce_builder import (
     make_ts_millis,
@@ -24,33 +27,77 @@ from datahub.emitter.mce_builder import (
 from datahub.emitter.mcp_builder import ContainerKey
 from datahub.errors import MultipleSubtypesWarning, SdkUsageError
 from datahub.metadata.urns import (
+    ChartUrn,
     ContainerUrn,
     CorpGroupUrn,
     CorpUserUrn,
+    DashboardUrn,
+    DataFlowUrn,
     DataJobUrn,
     DataPlatformInstanceUrn,
     DataPlatformUrn,
+    DataProcessInstanceUrn,
     DatasetUrn,
     DomainUrn,
     GlossaryTermUrn,
     OwnershipTypeUrn,
+    StructuredPropertyUrn,
     TagUrn,
     Urn,
+    VersionSetUrn,
 )
-from datahub.sdk._utils import add_list_unique, remove_list_unique
+from datahub.sdk._utils import DEFAULT_ACTOR_URN, add_list_unique, remove_list_unique
 from datahub.sdk.entity import Entity
 from datahub.utilities.urns.error import InvalidUrnError
 
 if TYPE_CHECKING:
     from datahub.sdk.container import Container
-
 UrnOrStr: TypeAlias = Union[Urn, str]
+ChartUrnOrStr: TypeAlias = Union[str, ChartUrn]
 DatasetUrnOrStr: TypeAlias = Union[str, DatasetUrn]
 DatajobUrnOrStr: TypeAlias = Union[str, DataJobUrn]
+DataflowUrnOrStr: TypeAlias = Union[str, DataFlowUrn]
+DashboardUrnOrStr: TypeAlias = Union[str, DashboardUrn]
+DataPlatformInstanceUrnOrStr: TypeAlias = Union[str, DataPlatformInstanceUrn]
+DataPlatformUrnOrStr: TypeAlias = Union[str, DataPlatformUrn]
 
 ActorUrn: TypeAlias = Union[CorpUserUrn, CorpGroupUrn]
+ActorUrnOrStr: TypeAlias = Union[str, ActorUrn]
+StructuredPropertyUrnOrStr: TypeAlias = Union[str, StructuredPropertyUrn]
+StructuredPropertyValueType: TypeAlias = Union[str, float, int]
+StructuredPropertyInputType: TypeAlias = Dict[
+    StructuredPropertyUrnOrStr, Sequence[StructuredPropertyValueType]
+]
+
+TrainingMetricsInputType: TypeAlias = Union[
+    List[models.MLMetricClass], Dict[str, Optional[str]]
+]
+HyperParamsInputType: TypeAlias = Union[
+    List[models.MLHyperParamClass], Dict[str, Optional[str]]
+]
+MLTrainingJobInputType: TypeAlias = Union[Sequence[Union[str, DataProcessInstanceUrn]]]
 
-_DEFAULT_ACTOR_URN = CorpUserUrn("__ingestion").urn()
+
+def convert_training_metrics(
+    metrics: TrainingMetricsInputType,
+) -> List[models.MLMetricClass]:
+    if isinstance(metrics, dict):
+        return [
+            models.MLMetricClass(name=name, value=str(value))
+            for name, value in metrics.items()
+        ]
+    return metrics
+
+
+def convert_hyper_params(
+    params: HyperParamsInputType,
+) -> List[models.MLHyperParamClass]:
+    if isinstance(params, dict):
+        return [
+            models.MLHyperParamClass(name=name, value=str(value))
+            for name, value in params.items()
+        ]
+    return params
 
 
 def make_time_stamp(ts: Optional[datetime]) -> Optional[models.TimeStampClass]:
@@ -65,6 +112,130 @@ def parse_time_stamp(ts: Optional[models.TimeStampClass]) -> Optional[datetime]:
     return parse_ts_millis(ts.time)
 
 
+class ChangeAuditStampsMixin(ABC):
+    """Mixin class for managing audit stamps on entities."""
+
+    __slots__ = ()
+
+    @abstractmethod
+    def _get_audit_stamps(self) -> models.ChangeAuditStampsClass:
+        """Get the audit stamps from the entity properties."""
+        pass
+
+    @abstractmethod
+    def _set_audit_stamps(self, audit_stamps: models.ChangeAuditStampsClass) -> None:
+        """Set the audit stamps on the entity properties."""
+        pass
+
+    @property
+    def last_modified(self) -> Optional[datetime]:
+        """Get the last modification timestamp from audit stamps."""
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        if audit_stamps.lastModified.time == 0:
+            return None
+        return datetime.fromtimestamp(
+            audit_stamps.lastModified.time / 1000
+        )  # supports only seconds precision
+
+    def set_last_modified(self, last_modified: datetime) -> None:
+        """Set the last modification timestamp in audit stamps."""
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        audit_stamps.lastModified.time = make_ts_millis(last_modified)
+        self._set_audit_stamps(audit_stamps)
+
+    @property
+    def last_modified_by(self) -> Optional[str]:
+        """Get the last modification actor from audit stamps."""
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        if audit_stamps.lastModified.actor == builder.UNKNOWN_USER:
+            return None
+        return audit_stamps.lastModified.actor
+
+    def set_last_modified_by(self, last_modified_by: ActorUrnOrStr) -> None:
+        """Set the last modification actor in audit stamps."""
+        if isinstance(last_modified_by, str):
+            last_modified_by = make_user_urn(last_modified_by)
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        audit_stamps.lastModified.actor = str(last_modified_by)
+        self._set_audit_stamps(audit_stamps)
+
+    @property
+    def created_at(self) -> Optional[datetime]:
+        """Get the creation timestamp from audit stamps."""
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        if audit_stamps.created.time == 0:
+            return None
+        return datetime.fromtimestamp(
+            audit_stamps.created.time / 1000
+        )  # supports only seconds precision
+
+    def set_created_at(self, created_at: datetime) -> None:
+        """Set the creation timestamp in audit stamps."""
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        audit_stamps.created.time = make_ts_millis(created_at)
+        self._set_audit_stamps(audit_stamps)
+
+    @property
+    def created_by(self) -> Optional[ActorUrnOrStr]:
+        """Get the creation actor from audit stamps."""
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        if audit_stamps.created.actor == builder.UNKNOWN_USER:
+            return None
+        return audit_stamps.created.actor
+
+    def set_created_by(self, created_by: ActorUrnOrStr) -> None:
+        """Set the creation actor in audit stamps."""
+        if isinstance(created_by, str):
+            created_by = make_user_urn(created_by)
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        audit_stamps.created.actor = str(created_by)
+        self._set_audit_stamps(audit_stamps)
+
+    @property
+    def deleted_on(self) -> Optional[datetime]:
+        """Get the deletion timestamp from audit stamps."""
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        if audit_stamps.deleted is None or audit_stamps.deleted.time == 0:
+            return None
+        return datetime.fromtimestamp(
+            audit_stamps.deleted.time / 1000
+        )  # supports only seconds precision
+
+    def set_deleted_on(self, deleted_on: datetime) -> None:
+        """Set the deletion timestamp in audit stamps."""
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        # Default constructor sets deleted to None
+        if audit_stamps.deleted is None:
+            audit_stamps.deleted = models.AuditStampClass(
+                time=0, actor=builder.UNKNOWN_USER
+            )
+        audit_stamps.deleted.time = make_ts_millis(deleted_on)
+        self._set_audit_stamps(audit_stamps)
+
+    @property
+    def deleted_by(self) -> Optional[ActorUrnOrStr]:
+        """Get the deletion actor from audit stamps."""
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        if (
+            audit_stamps.deleted is None
+            or audit_stamps.deleted.actor == builder.UNKNOWN_USER
+        ):
+            return None
+        return audit_stamps.deleted.actor
+
+    def set_deleted_by(self, deleted_by: ActorUrnOrStr) -> None:
+        """Set the deletion actor in audit stamps."""
+        if isinstance(deleted_by, str):
+            deleted_by = make_user_urn(deleted_by)
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        if audit_stamps.deleted is None:
+            audit_stamps.deleted = models.AuditStampClass(
+                time=0, actor=builder.UNKNOWN_USER
+            )
+        audit_stamps.deleted.actor = str(deleted_by)
+        self._set_audit_stamps(audit_stamps)
+
+
 class HasPlatformInstance(Entity):
     __slots__ = ()
 
@@ -134,7 +305,7 @@ OwnerInputType: TypeAlias = Union[
     Tuple[ActorUrn, OwnershipTypeType],
     models.OwnerClass,
 ]
-OwnersInputType: TypeAlias = List[OwnerInputType]
+OwnersInputType: TypeAlias = Sequence[OwnerInputType]
 
 
 class HasOwnership(Entity):
@@ -235,7 +406,9 @@ class HasOwnership(Entity):
 # If you pass in a ContainerKey, we can use parent_key() to build the browse path.
 # If you pass in a list of urns, we'll use that as the browse path. Any non-urn strings
 # will be treated as raw ids.
-ParentContainerInputType: TypeAlias = Union["Container", ContainerKey, List[UrnOrStr]]
+ParentContainerInputType: TypeAlias = Union[
+    "Container", ContainerKey, Sequence[UrnOrStr]
+]
 
 
 class HasContainer(Entity):
@@ -295,7 +468,7 @@ class HasContainer(Entity):
                 )
                 for entry in parsed_path
             ]
-        elif container is not None:
+        elif isinstance(container, ContainerKey):
             container_urn = container.as_urn()
 
             browse_path_reversed = [container_urn]
@@ -354,7 +527,7 @@ class HasContainer(Entity):
 
 
 TagInputType: TypeAlias = Union[str, TagUrn, models.TagAssociationClass]
-TagsInputType: TypeAlias = List[TagInputType]
+TagsInputType: TypeAlias = Sequence[TagInputType]
 
 
 class HasTags(Entity):
@@ -409,7 +582,7 @@ class HasTags(Entity):
 TermInputType: TypeAlias = Union[
     str, GlossaryTermUrn, models.GlossaryTermAssociationClass
 ]
-TermsInputType: TypeAlias = List[TermInputType]
+TermsInputType: TypeAlias = Sequence[TermInputType]
 
 
 class HasTerms(Entity):
@@ -441,7 +614,7 @@ class HasTerms(Entity):
     def _terms_audit_stamp(self) -> models.AuditStampClass:
         return models.AuditStampClass(
             time=0,
-            actor=_DEFAULT_ACTOR_URN,
+            actor=DEFAULT_ACTOR_URN,
         )
 
     def set_terms(self, terms: TermsInputType) -> None:
@@ -529,7 +702,7 @@ class HasInstitutionalMemory(Entity):
     def _institutional_memory_audit_stamp(self) -> models.AuditStampClass:
         return models.AuditStampClass(
             time=0,
-            actor=_DEFAULT_ACTOR_URN,
+            actor=DEFAULT_ACTOR_URN,
         )
 
     @classmethod
@@ -578,3 +751,213 @@ class HasInstitutionalMemory(Entity):
             self._link_key,
             self._parse_link_association_class(link),
         )
+
+
+class HasVersion(Entity):
+    """Mixin for entities that have version properties."""
+
+    def _get_version_props(self) -> Optional[models.VersionPropertiesClass]:
+        return self._get_aspect(models.VersionPropertiesClass)
+
+    def _ensure_version_props(self) -> models.VersionPropertiesClass:
+        version_props = self._get_version_props()
+        if version_props is None:
+            guid_dict = {"urn": str(self.urn)}
+            version_set_urn = VersionSetUrn(
+                id=builder.datahub_guid(guid_dict), entity_type=self.urn.ENTITY_TYPE
+            )
+
+            version_props = models.VersionPropertiesClass(
+                versionSet=str(version_set_urn),
+                version=models.VersionTagClass(versionTag="0.1.0"),
+                sortId="0000000.1.0",
+            )
+            self._set_aspect(version_props)
+        return version_props
+
+    @property
+    def version(self) -> Optional[str]:
+        version_props = self._get_version_props()
+        if version_props and version_props.version:
+            return version_props.version.versionTag
+        return None
+
+    def set_version(self, version: str) -> None:
+        """Set the version of the entity."""
+        guid_dict = {"urn": str(self.urn)}
+        version_set_urn = VersionSetUrn(
+            id=builder.datahub_guid(guid_dict), entity_type=self.urn.ENTITY_TYPE
+        )
+
+        version_props = self._get_version_props()
+        if version_props is None:
+            # If no version properties exist, create a new one
+            version_props = models.VersionPropertiesClass(
+                version=models.VersionTagClass(versionTag=version),
+                versionSet=str(version_set_urn),
+                sortId=version.zfill(10),  # Pad with zeros for sorting
+            )
+        else:
+            # Update existing version properties
+            version_props.version = models.VersionTagClass(versionTag=version)
+            version_props.versionSet = str(version_set_urn)
+            version_props.sortId = version.zfill(10)
+
+        self._set_aspect(version_props)
+
+    @property
+    def version_aliases(self) -> List[str]:
+        version_props = self._get_version_props()
+        if version_props and version_props.aliases:
+            return [
+                alias.versionTag
+                for alias in version_props.aliases
+                if alias.versionTag is not None
+            ]
+        return []  # Return empty list instead of None
+
+    def set_version_aliases(self, aliases: List[str]) -> None:
+        version_props = self._get_aspect(models.VersionPropertiesClass)
+        if version_props:
+            version_props.aliases = [
+                models.VersionTagClass(versionTag=alias) for alias in aliases
+            ]
+        else:
+            # If no version properties exist, we need to create one with a default version
+            guid_dict = {"urn": str(self.urn)}
+            version_set_urn = VersionSetUrn(
+                id=builder.datahub_guid(guid_dict), entity_type=self.urn.ENTITY_TYPE
+            )
+            self._set_aspect(
+                models.VersionPropertiesClass(
+                    version=models.VersionTagClass(
+                        versionTag="0.1.0"
+                    ),  # Default version
+                    versionSet=str(version_set_urn),
+                    sortId="0000000.1.0",
+                    aliases=[
+                        models.VersionTagClass(versionTag=alias) for alias in aliases
+                    ],
+                )
+            )
+
+    def add_version_alias(self, alias: str) -> None:
+        if not alias:
+            raise ValueError("Alias cannot be empty")
+        version_props = self._ensure_version_props()
+        if version_props.aliases is None:
+            version_props.aliases = []
+        version_props.aliases.append(models.VersionTagClass(versionTag=alias))
+        self._set_aspect(version_props)
+
+    def remove_version_alias(self, alias: str) -> None:
+        version_props = self._get_version_props()
+        if version_props and version_props.aliases:
+            version_props.aliases = [
+                a for a in version_props.aliases if a.versionTag != alias
+            ]
+            self._set_aspect(version_props)
+
+
+class HasStructuredProperties(Entity):
+    """
+    Mixin for entities that support structured properties
+    """
+
+    __slots__ = ()
+
+    @property
+    def structured_properties(
+        self,
+    ) -> Optional[List[models.StructuredPropertyValueAssignmentClass]]:
+        """
+        Retrieve structured properties for the entity
+
+        Returns:
+            Optional list of structured property value assignments
+        """
+        sp_aspect = self._get_aspect(models.StructuredPropertiesClass)
+        return sp_aspect.properties if sp_aspect else None
+
+    def _ensure_structured_properties(self) -> models.StructuredPropertiesClass:
+        """
+        Ensure structured properties aspect exists, creating it if necessary
+
+        Returns:
+            StructuredPropertiesClass aspect
+        """
+        return self._setdefault_aspect(models.StructuredPropertiesClass(properties=[]))
+
+    def set_structured_property(
+        self,
+        property_urn: StructuredPropertyUrnOrStr,
+        values: Sequence[StructuredPropertyValueType],
+    ) -> None:
+        """
+        Update an existing structured property or add if it doesn't exist
+
+        Args:
+            property_urn: URN of the structured property
+            values: List of values for the property
+        """
+        # validate property_urn is a valid structured property urn
+        property_urn = StructuredPropertyUrn.from_string(property_urn)
+
+        properties = self._ensure_structured_properties()
+
+        # Find existing property assignment
+        existing_prop = next(
+            (
+                prop
+                for prop in properties.properties
+                if prop.propertyUrn == str(property_urn)
+            ),
+            None,
+        )
+        current_timestamp = make_ts_millis(datetime.now())
+
+        if existing_prop:
+            # Update existing property
+            existing_prop.values = list(values)
+            existing_prop.lastModified = models.AuditStampClass(
+                time=current_timestamp,
+                actor=DEFAULT_ACTOR_URN,
+            )
+        else:
+            # Create new property assignment
+            new_property = models.StructuredPropertyValueAssignmentClass(
+                propertyUrn=str(property_urn),
+                values=list(values),
+                created=models.AuditStampClass(
+                    time=current_timestamp,
+                    actor=DEFAULT_ACTOR_URN,
+                ),
+                lastModified=models.AuditStampClass(
+                    time=current_timestamp,
+                    actor=DEFAULT_ACTOR_URN,
+                ),
+            )
+            add_list_unique(
+                properties.properties,
+                key=lambda prop: prop.propertyUrn,
+                item=new_property,
+            )
+
+        self._set_aspect(properties)
+
+    def remove_structured_property(
+        self, property_urn: StructuredPropertyUrnOrStr
+    ) -> None:
+        """
+        Remove a structured property from the entity
+
+        Args:
+            property_urn: URN of the structured property to remove
+        """
+        remove_list_unique(
+            self._ensure_structured_properties().properties,
+            key=lambda prop: prop.propertyUrn,
+            item=models.StructuredPropertyValueAssignmentClass(
+                propertyUrn=str(property_urn), values=[]
+            ),
+        )

datahub/sdk/_utils.py

@@ -1,6 +1,10 @@
 from typing import Any, Callable, List, Protocol, TypeVar
 
 from datahub.errors import ItemNotFoundError
+from datahub.metadata.urns import CorpUserUrn
+
+# TODO: Change __ingestion to _ingestion.
+DEFAULT_ACTOR_URN = CorpUserUrn("__ingestion").urn()
 
 
 class _SupportsEq(Protocol):