acryl-datahub 0.15.0.5rc10__py3-none-any.whl → 0.15.0.6rc2__py3-none-any.whl

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
+
+    @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)