acryl-datahub 0.15.0.6rc1__py3-none-any.whl → 0.15.0.6rc3__py3-none-any.whl

datahub/sdk/_all_entities.py

@@ -0,0 +1,15 @@
+from typing import Dict, List, Type
+
+from datahub.sdk._entity import Entity
+from datahub.sdk.container import Container
+from datahub.sdk.dataset import Dataset
+
+# TODO: Is there a better way to declare this?
+ENTITY_CLASSES_LIST: List[Type[Entity]] = [
+    Container,
+    Dataset,
+]
+
+ENTITY_CLASSES: Dict[str, Type[Entity]] = {
+    cls.get_urn_type().ENTITY_TYPE: cls for cls in ENTITY_CLASSES_LIST
+}

datahub/sdk/_attribution.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import contextlib
+from typing import Iterator
+
+from datahub.utilities.str_enum import StrEnum
+
+
+class KnownAttribution(StrEnum):
+    INGESTION = "INGESTION"
+    INGESTION_ALTERNATE = "INGESTION_ALTERNATE"
+
+    UI = "UI"
+    SDK = "SDK"
+
+    PROPAGATION = "PROPAGATION"
+
+    def is_ingestion(self) -> bool:
+        return self in (
+            KnownAttribution.INGESTION,
+            KnownAttribution.INGESTION_ALTERNATE,
+        )
+
+
+_default_attribution = KnownAttribution.SDK
+
+
+def get_default_attribution() -> KnownAttribution:
+    return _default_attribution
+
+
+def set_default_attribution(attribution: KnownAttribution) -> None:
+    global _default_attribution
+    _default_attribution = attribution
+
+
+@contextlib.contextmanager
+def change_default_attribution(attribution: KnownAttribution) -> Iterator[None]:
+    old_attribution = get_default_attribution()
+    try:
+        set_default_attribution(attribution)
+        yield
+    finally:
+        set_default_attribution(old_attribution)
+
+
+def is_ingestion_attribution() -> bool:
+    return get_default_attribution().is_ingestion()

datahub/sdk/_entity.py

@@ -0,0 +1,89 @@
+import abc
+from typing import List, Optional, Type, Union
+
+from typing_extensions import Self
+
+import datahub.metadata.schema_classes as models
+from datahub.emitter.mce_builder import Aspect as AspectTypeVar
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.errors import SdkUsageError
+from datahub.metadata.urns import Urn
+from datahub.utilities.urns._urn_base import _SpecificUrn
+
+
+class Entity:
+    __slots__ = ("_urn", "_prev_aspects", "_aspects")
+
+    def __init__(self, /, urn: Urn):
+        # This method is not meant for direct usage.
+        if type(self) is Entity:
+            raise SdkUsageError(f"{Entity.__name__} cannot be instantiated directly.")
+
+        assert isinstance(urn, self.get_urn_type())
+        self._urn: _SpecificUrn = urn
+
+        # prev_aspects is None means this was created from scratch
+        self._prev_aspects: Optional[models.AspectBag] = None
+        self._aspects: models.AspectBag = {}
+
+    @classmethod
+    def _new_from_graph(cls, urn: Urn, current_aspects: models.AspectBag) -> Self:
+        # If an init method from a subclass adds required fields, it also needs to override this method.
+        # An alternative approach would call cls.__new__() to bypass the init method, but it's a bit
+        # too hacky for my taste.
+        entity = cls(urn=urn)
+        return entity._init_from_graph(current_aspects)
+
+    def _init_from_graph(self, current_aspects: models.AspectBag) -> Self:
+        self._prev_aspects = current_aspects
+        aspect: models._Aspect
+        for aspect_name, aspect in (current_aspects or {}).items():  # type: ignore
+            aspect_copy = type(aspect).from_obj(aspect.to_obj())
+            self._aspects[aspect_name] = aspect_copy  # type: ignore
+        return self
+
+    @classmethod
+    @abc.abstractmethod
+    def get_urn_type(cls) -> Type[_SpecificUrn]: ...
+
+    @property
+    def urn(self) -> _SpecificUrn:
+        return self._urn
+
+    def _get_aspect(
+        self,
+        aspect_type: Type[AspectTypeVar],
+        /,
+    ) -> Optional[AspectTypeVar]:
+        return self._aspects.get(aspect_type.ASPECT_NAME)  # type: ignore
+
+    def _set_aspect(self, value: AspectTypeVar, /) -> None:
+        self._aspects[value.ASPECT_NAME] = value  # type: ignore
+
+    def _setdefault_aspect(self, default_aspect: AspectTypeVar, /) -> AspectTypeVar:
+        # Similar semantics to dict.setdefault.
+        if existing_aspect := self._get_aspect(type(default_aspect)):
+            return existing_aspect
+        self._set_aspect(default_aspect)
+        return default_aspect
+
+    def _as_mcps(
+        self,
+        change_type: Union[str, models.ChangeTypeClass] = models.ChangeTypeClass.UPSERT,
+    ) -> List[MetadataChangeProposalWrapper]:
+        urn_str = str(self.urn)
+
+        mcps = []
+        for aspect in self._aspects.values():
+            assert isinstance(aspect, models._Aspect)
+            mcps.append(
+                MetadataChangeProposalWrapper(
+                    entityUrn=urn_str,
+                    aspect=aspect,
+                    changeType=change_type,
+                )
+            )
+        return mcps
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}('{self.urn}')"

datahub/sdk/_shared.py

@@ -0,0 +1,345 @@
+import warnings
+from datetime import datetime
+from typing import (
+    TYPE_CHECKING,
+    List,
+    Optional,
+    Tuple,
+    Union,
+)
+
+from typing_extensions import TypeAlias
+
+import datahub.metadata.schema_classes as models
+from datahub.emitter.mce_builder import (
+    make_ts_millis,
+    make_user_urn,
+    parse_ts_millis,
+    validate_ownership_type,
+)
+from datahub.emitter.mcp_builder import ContainerKey
+from datahub.errors import MultipleSubtypesWarning, SdkUsageError
+from datahub.metadata.urns import (
+    CorpGroupUrn,
+    CorpUserUrn,
+    DataJobUrn,
+    DataPlatformInstanceUrn,
+    DataPlatformUrn,
+    DatasetUrn,
+    DomainUrn,
+    GlossaryTermUrn,
+    OwnershipTypeUrn,
+    TagUrn,
+    Urn,
+)
+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]
+DatasetUrnOrStr: TypeAlias = Union[str, DatasetUrn]
+DatajobUrnOrStr: TypeAlias = Union[str, DataJobUrn]
+
+ActorUrn: TypeAlias = Union[CorpUserUrn, CorpGroupUrn]
+
+
+def make_time_stamp(ts: Optional[datetime]) -> Optional[models.TimeStampClass]:
+    if ts is None:
+        return None
+    return models.TimeStampClass(time=make_ts_millis(ts))
+
+
+def parse_time_stamp(ts: Optional[models.TimeStampClass]) -> Optional[datetime]:
+    if ts is None:
+        return None
+    return parse_ts_millis(ts.time)
+
+
+class HasPlatformInstance(Entity):
+    __slots__ = ()
+
+    def _set_platform_instance(
+        self,
+        platform: Union[str, DataPlatformUrn],
+        instance: Union[None, str, DataPlatformInstanceUrn],
+    ) -> None:
+        platform = DataPlatformUrn(platform)
+        if instance is not None:
+            try:
+                instance = DataPlatformInstanceUrn.from_string(instance)
+            except InvalidUrnError:
+                if not isinstance(
+                    instance, DataPlatformInstanceUrn
+                ):  # redundant check to make mypy happy
+                    instance = DataPlatformInstanceUrn(platform, instance)
+        # At this point, instance is either None or a DataPlatformInstanceUrn.
+
+        self._set_aspect(
+            models.DataPlatformInstanceClass(
+                platform=platform.urn(),
+                instance=instance.urn() if instance else None,
+            )
+        )
+
+    @property
+    def platform_instance(self) -> Optional[DataPlatformInstanceUrn]:
+        dataPlatformInstance = self._get_aspect(models.DataPlatformInstanceClass)
+        if dataPlatformInstance and dataPlatformInstance.instance:
+            return DataPlatformInstanceUrn.from_string(dataPlatformInstance.instance)
+        return None
+
+
+class HasSubtype(Entity):
+    __slots__ = ()
+
+    @property
+    def subtype(self) -> Optional[str]:
+        subtypes = self._get_aspect(models.SubTypesClass)
+        if subtypes and subtypes.typeNames:
+            if len(subtypes.typeNames) > 1:
+                warnings.warn(
+                    f"The entity {self.urn} has multiple subtypes: {subtypes.typeNames}. "
+                    "Only the first subtype will be considered.",
+                    MultipleSubtypesWarning,
+                    stacklevel=2,
+                )
+            return subtypes.typeNames[0]
+        return None
+
+    def set_subtype(self, subtype: str) -> None:
+        self._set_aspect(models.SubTypesClass(typeNames=[subtype]))
+
+
+OwnershipTypeType: TypeAlias = Union[str, OwnershipTypeUrn]
+OwnerInputType: TypeAlias = Union[
+    str,
+    ActorUrn,
+    Tuple[Union[str, ActorUrn], OwnershipTypeType],
+    models.OwnerClass,
+]
+OwnersInputType: TypeAlias = List[OwnerInputType]
+
+
+class HasOwnership(Entity):
+    __slots__ = ()
+
+    @staticmethod
+    def _parse_owner_class(owner: OwnerInputType) -> models.OwnerClass:
+        if isinstance(owner, models.OwnerClass):
+            return owner
+
+        owner_type = models.OwnershipTypeClass.TECHNICAL_OWNER
+        owner_type_urn = None
+
+        if isinstance(owner, tuple):
+            raw_owner, raw_owner_type = owner
+
+            if isinstance(raw_owner_type, OwnershipTypeUrn):
+                owner_type = models.OwnershipTypeClass.CUSTOM
+                owner_type_urn = str(raw_owner_type)
+            else:
+                owner_type, owner_type_urn = validate_ownership_type(raw_owner_type)
+        else:
+            raw_owner = owner
+
+        if isinstance(raw_owner, str):
+            # Tricky: this will gracefully handle a user passing in a group urn as a string.
+            # TODO: is this the right behavior? or should we require a valid urn here?
+            return models.OwnerClass(
+                owner=make_user_urn(raw_owner),
+                type=owner_type,
+                typeUrn=owner_type_urn,
+            )
+        elif isinstance(raw_owner, Urn):
+            return models.OwnerClass(
+                owner=str(raw_owner),
+                type=owner_type,
+                typeUrn=owner_type_urn,
+            )
+        else:
+            raise SdkUsageError(
+                f"Invalid owner {owner}: {type(owner)} is not a valid owner type"
+            )
+
+    # TODO: Return a custom type with deserialized urns, instead of the raw aspect.
+    # Ideally we'd also use first-class ownership type urns here, not strings.
+    @property
+    def owners(self) -> Optional[List[models.OwnerClass]]:
+        if owners_aspect := self._get_aspect(models.OwnershipClass):
+            return owners_aspect.owners
+        return None
+
+    def set_owners(self, owners: OwnersInputType) -> None:
+        # TODO: add docs on the default parsing + default ownership type
+        parsed_owners = [self._parse_owner_class(owner) for owner in owners]
+        self._set_aspect(models.OwnershipClass(owners=parsed_owners))
+
+
+ContainerInputType: TypeAlias = Union["Container", ContainerKey]
+
+
+class HasContainer(Entity):
+    __slots__ = ()
+
+    def _set_container(self, container: Optional[ContainerInputType]) -> None:
+        # We need to allow container to be None. It won't happen for datasets much, but
+        # will be required for root containers.
+        from datahub.sdk.container import Container
+
+        browse_path: List[Union[str, models.BrowsePathEntryClass]] = []
+        if isinstance(container, Container):
+            container_urn = container.urn.urn()
+
+            parent_browse_path = container._get_aspect(models.BrowsePathsV2Class)
+            if parent_browse_path is None:
+                raise SdkUsageError(
+                    "Parent container does not have a browse path, so cannot generate one for its children."
+                )
+            browse_path = [
+                *parent_browse_path.path,
+                models.BrowsePathEntryClass(
+                    id=container_urn,
+                    urn=container_urn,
+                ),
+            ]
+        elif container is not None:
+            container_urn = container.as_urn()
+
+            browse_path_reversed = [container_urn]
+            parent_key = container.parent_key()
+            while parent_key is not None:
+                browse_path_reversed.append(parent_key.as_urn())
+                parent_key = parent_key.parent_key()
+            if container.instance is not None:
+                browse_path_reversed.append(
+                    DataPlatformInstanceUrn(
+                        container.platform, container.instance
+                    ).urn()
+                )
+
+            browse_path = list(reversed(browse_path_reversed))
+        else:
+            container_urn = None
+            browse_path = []
+
+        if container_urn:
+            self._set_aspect(models.ContainerClass(container=container_urn))
+
+        self._set_aspect(
+            models.BrowsePathsV2Class(
+                path=[
+                    (
+                        entry
+                        if isinstance(entry, models.BrowsePathEntryClass)
+                        else models.BrowsePathEntryClass(
+                            id=entry,
+                            urn=entry,
+                        )
+                    )
+                    for entry in browse_path
+                ]
+            )
+        )
+
+
+TagInputType: TypeAlias = Union[str, TagUrn, models.TagAssociationClass]
+TagsInputType: TypeAlias = List[TagInputType]
+
+
+class HasTags(Entity):
+    __slots__ = ()
+
+    # TODO: Return a custom type with deserialized urns, instead of the raw aspect.
+    @property
+    def tags(self) -> Optional[List[models.TagAssociationClass]]:
+        if tags := self._get_aspect(models.GlobalTagsClass):
+            return tags.tags
+        return None
+
+    @classmethod
+    def _parse_tag_association_class(
+        cls, tag: TagInputType
+    ) -> models.TagAssociationClass:
+        if isinstance(tag, models.TagAssociationClass):
+            return tag
+        elif isinstance(tag, str):
+            assert TagUrn.from_string(tag)
+        return models.TagAssociationClass(tag=str(tag))
+
+    def set_tags(self, tags: TagsInputType) -> None:
+        self._set_aspect(
+            models.GlobalTagsClass(
+                tags=[self._parse_tag_association_class(tag) for tag in tags]
+            )
+        )
+
+
+TermInputType: TypeAlias = Union[
+    str, GlossaryTermUrn, models.GlossaryTermAssociationClass
+]
+TermsInputType: TypeAlias = List[TermInputType]
+
+
+class HasTerms(Entity):
+    __slots__ = ()
+
+    # TODO: Return a custom type with deserialized urns, instead of the raw aspect.
+    @property
+    def terms(self) -> Optional[List[models.GlossaryTermAssociationClass]]:
+        if glossary_terms := self._get_aspect(models.GlossaryTermsClass):
+            return glossary_terms.terms
+        return None
+
+    @classmethod
+    def _parse_glossary_term_association_class(
+        cls, term: TermInputType
+    ) -> models.GlossaryTermAssociationClass:
+        if isinstance(term, models.GlossaryTermAssociationClass):
+            return term
+        elif isinstance(term, str):
+            assert GlossaryTermUrn.from_string(term)
+        return models.GlossaryTermAssociationClass(urn=str(term))
+
+    @classmethod
+    def _terms_audit_stamp(self) -> models.AuditStampClass:
+        return models.AuditStampClass(
+            time=0,
+            # TODO figure out what to put here
+            actor=CorpUserUrn("__ingestion").urn(),
+        )
+
+    def set_terms(self, terms: TermsInputType) -> None:
+        self._set_aspect(
+            models.GlossaryTermsClass(
+                terms=[
+                    self._parse_glossary_term_association_class(term) for term in terms
+                ],
+                auditStamp=self._terms_audit_stamp(),
+            )
+        )
+
+
+DomainInputType: TypeAlias = Union[str, DomainUrn]
+
+
+class HasDomain(Entity):
+    __slots__ = ()
+
+    @property
+    def domain(self) -> Optional[DomainUrn]:
+        if domains := self._get_aspect(models.DomainsClass):
+            if len(domains.domains) > 1:
+                raise SdkUsageError(
+                    f"The entity has multiple domains set, but only one is supported: {domains.domains}"
+                )
+            elif domains.domains:
+                domain_str = domains.domains[0]
+                return DomainUrn.from_string(domain_str)
+
+        return None
+
+    def set_domain(self, domain: DomainInputType) -> None:
+        domain_urn = DomainUrn.from_string(domain)  # basically a type assertion
+        self._set_aspect(models.DomainsClass(domains=[str(domain_urn)]))

datahub/sdk/container.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Dict, Optional, Type
+
+from typing_extensions import Self
+
+import datahub.metadata.schema_classes as models
+from datahub.emitter.mce_builder import ALL_ENV_TYPES
+from datahub.emitter.mcp_builder import (
+    _INCLUDE_ENV_IN_CONTAINER_PROPERTIES,
+    ContainerKey,
+)
+from datahub.errors import SdkUsageError
+from datahub.metadata.urns import (
+    ContainerUrn,
+    Urn,
+)
+from datahub.sdk._entity import Entity
+from datahub.sdk._shared import (
+    DomainInputType,
+    HasContainer,
+    HasDomain,
+    HasOwnership,
+    HasPlatformInstance,
+    HasSubtype,
+    HasTags,
+    HasTerms,
+    OwnersInputType,
+    TagsInputType,
+    TermsInputType,
+    make_time_stamp,
+    parse_time_stamp,
+)
+
+
+class Container(
+    HasPlatformInstance,
+    HasSubtype,
+    HasContainer,
+    HasOwnership,
+    HasTags,
+    HasTerms,
+    HasDomain,
+    Entity,
+):
+    __slots__ = ()
+
+    @classmethod
+    def get_urn_type(cls) -> Type[ContainerUrn]:
+        return ContainerUrn
+
+    def __init__(
+        self,
+        /,
+        # Identity.
+        container_key: ContainerKey | ContainerUrn,
+        *,
+        # Container attributes.
+        display_name: str,
+        qualified_name: Optional[str] = None,
+        description: Optional[str] = None,
+        external_url: Optional[str] = None,
+        # TODO: call this custom properties?
+        extra_properties: Optional[Dict[str, str]] = None,
+        created: Optional[datetime] = None,
+        last_modified: Optional[datetime] = None,
+        # Standard aspects.
+        subtype: Optional[str] = None,
+        owners: Optional[OwnersInputType] = None,
+        tags: Optional[TagsInputType] = None,
+        terms: Optional[TermsInputType] = None,
+        domain: Optional[DomainInputType] = None,
+    ):
+        if isinstance(container_key, ContainerUrn):
+            urn = container_key
+        else:
+            urn = ContainerUrn.from_string(container_key.as_urn())
+        super().__init__(urn)
+
+        # This needs to come first to ensure that the display name is registered.
+        self._ensure_container_props(name=display_name)
+
+        # TODO: Normal usages should require container key. Only the graph init method can accept an urn.
+        if isinstance(container_key, ContainerKey):
+            self._set_platform_instance(container_key.platform, container_key.instance)
+
+            self._set_container(container_key.parent_key())
+
+            self.set_custom_properties(
+                {
+                    **container_key.property_dict(),
+                    **(extra_properties or {}),
+                }
+            )
+
+            # Extra validation on the env field.
+            # In certain cases (mainly for backwards compatibility), the env field will actually
+            # have a platform instance name.
+            env = container_key.env if container_key.env in ALL_ENV_TYPES else None
+            if _INCLUDE_ENV_IN_CONTAINER_PROPERTIES and env is not None:
+                self._ensure_container_props().env = env
+
+        if description is not None:
+            self.set_description(description)
+        if external_url is not None:
+            self.set_external_url(external_url)
+        if qualified_name is not None:
+            self.set_qualified_name(qualified_name)
+        if created is not None:
+            self.set_created(created)
+        if last_modified is not None:
+            self.set_last_modified(last_modified)
+
+        if subtype is not None:
+            self.set_subtype(subtype)
+        if owners is not None:
+            self.set_owners(owners)
+        if tags is not None:
+            self.set_tags(tags)
+        if terms is not None:
+            self.set_terms(terms)
+        if domain is not None:
+            self.set_domain(domain)
+
+    @classmethod
+    def _new_from_graph(cls, urn: Urn, current_aspects: models.AspectBag) -> Self:
+        assert isinstance(urn, ContainerUrn)
+        entity = cls(urn, display_name="__dummy_value__")
+        return entity._init_from_graph(current_aspects)
+
+    def _ensure_container_props(
+        self, *, name: Optional[str] = None
+    ) -> models.ContainerPropertiesClass:
+        # TODO: Not super happy with this method's implementation, but it's
+        # internal-only and enforces the constraints that we need.
+        if name is not None:
+            return self._setdefault_aspect(models.ContainerPropertiesClass(name=name))
+
+        props = self._get_aspect(models.ContainerPropertiesClass)
+        if props is None:
+            raise SdkUsageError("Containers must have a name.")
+        return props
+
+    @property
+    def display_name(self) -> str:
+        return self._ensure_container_props().name
+
+    def set_display_name(self, value: str) -> None:
+        self._ensure_container_props(name=value).name = value
+
+    @property
+    def description(self) -> Optional[str]:
+        return self._ensure_container_props().description
+
+    def set_description(self, description: str) -> None:
+        self._ensure_container_props().description = description
+
+    @property
+    def custom_properties(self) -> Optional[Dict[str, str]]:
+        return self._ensure_container_props().customProperties
+
+    def set_custom_properties(self, custom_properties: Dict[str, str]) -> None:
+        # TODO: How do we ensure that the container key props are always retained?
+        self._ensure_container_props().customProperties = custom_properties
+
+    @property
+    def external_url(self) -> Optional[str]:
+        return self._ensure_container_props().externalUrl
+
+    def set_external_url(self, external_url: str) -> None:
+        self._ensure_container_props().externalUrl = external_url
+
+    @property
+    def qualified_name(self) -> Optional[str]:
+        return self._ensure_container_props().qualifiedName
+
+    def set_qualified_name(self, qualified_name: str) -> None:
+        self._ensure_container_props().qualifiedName = qualified_name
+
+    @property
+    def created(self) -> Optional[datetime]:
+        return parse_time_stamp(self._ensure_container_props().created)
+
+    def set_created(self, created: datetime) -> None:
+        self._ensure_container_props().created = make_time_stamp(created)
+
+    @property
+    def last_modified(self) -> Optional[datetime]:
+        return parse_time_stamp(self._ensure_container_props().lastModified)
+
+    def set_last_modified(self, last_modified: datetime) -> None:
+        self._ensure_container_props().lastModified = make_time_stamp(last_modified)