acryl-datahub 1.0.0rc4__py3-none-any.whl → 1.0.0rc6__py3-none-any.whl

datahub/sdk/_shared.py

@@ -1,14 +1,17 @@
+from __future__ import annotations
+
 import warnings
 from datetime import datetime
 from typing import (
     TYPE_CHECKING,
+    Callable,
     List,
     Optional,
     Tuple,
     Union,
 )
 
-from typing_extensions import TypeAlias
+from typing_extensions import TypeAlias, assert_never
 
 import datahub.metadata.schema_classes as models
 from datahub.emitter.mce_builder import (
@@ -20,6 +23,7 @@ from datahub.emitter.mce_builder import (
 from datahub.emitter.mcp_builder import ContainerKey
 from datahub.errors import MultipleSubtypesWarning, SdkUsageError
 from datahub.metadata.urns import (
+    ContainerUrn,
     CorpGroupUrn,
     CorpUserUrn,
     DataJobUrn,
@@ -33,6 +37,7 @@ from datahub.metadata.urns import (
     Urn,
 )
 from datahub.sdk._entity import Entity
+from datahub.sdk._utils import add_list_unique, remove_list_unique
 from datahub.utilities.urns.error import InvalidUrnError
 
 if TYPE_CHECKING:
@@ -83,6 +88,13 @@ class HasPlatformInstance(Entity):
             )
         )
 
+    @property
+    def platform(self) -> Optional[DataPlatformUrn]:
+        dataPlatform = self._get_aspect(models.DataPlatformInstanceClass)
+        if dataPlatform and dataPlatform.platform:
+            return DataPlatformUrn.from_string(dataPlatform.platform)
+        return None
+
     @property
     def platform_instance(self) -> Optional[DataPlatformInstanceUrn]:
         dataPlatformInstance = self._get_aspect(models.DataPlatformInstanceClass)
@@ -112,11 +124,11 @@ class HasSubtype(Entity):
         self._set_aspect(models.SubTypesClass(typeNames=[subtype]))
 
 
+# TODO: Reference OwnershipTypeClass as the valid ownership type enum.
 OwnershipTypeType: TypeAlias = Union[str, OwnershipTypeUrn]
 OwnerInputType: TypeAlias = Union[
-    str,
     ActorUrn,
-    Tuple[Union[str, ActorUrn], OwnershipTypeType],
+    Tuple[ActorUrn, OwnershipTypeType],
     models.OwnerClass,
 ]
 OwnersInputType: TypeAlias = List[OwnerInputType]
@@ -126,15 +138,17 @@ class HasOwnership(Entity):
     __slots__ = ()
 
     @staticmethod
-    def _parse_owner_class(owner: OwnerInputType) -> models.OwnerClass:
+    def _parse_owner_class(owner: OwnerInputType) -> Tuple[models.OwnerClass, bool]:
         if isinstance(owner, models.OwnerClass):
-            return owner
+            return owner, False
 
+        was_type_specified = False
         owner_type = models.OwnershipTypeClass.TECHNICAL_OWNER
         owner_type_urn = None
 
         if isinstance(owner, tuple):
             raw_owner, raw_owner_type = owner
+            was_type_specified = True
 
             if isinstance(raw_owner_type, OwnershipTypeUrn):
                 owner_type = models.OwnershipTypeClass.CUSTOM
@@ -151,17 +165,15 @@ class HasOwnership(Entity):
                 owner=make_user_urn(raw_owner),
                 type=owner_type,
                 typeUrn=owner_type_urn,
-            )
+            ), was_type_specified
         elif isinstance(raw_owner, Urn):
             return models.OwnerClass(
                 owner=str(raw_owner),
                 type=owner_type,
                 typeUrn=owner_type_urn,
-            )
+            ), was_type_specified
         else:
-            raise SdkUsageError(
-                f"Invalid owner {owner}: {type(owner)} is not a valid owner type"
-            )
+            assert_never(raw_owner)
 
     # 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.
@@ -173,21 +185,74 @@ class HasOwnership(Entity):
 
     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]
+        parsed_owners = [self._parse_owner_class(owner)[0] for owner in owners]
         self._set_aspect(models.OwnershipClass(owners=parsed_owners))
 
+    @classmethod
+    def _owner_key_method(
+        cls, consider_owner_type: bool
+    ) -> Callable[[models.OwnerClass], Tuple[str, ...]]:
+        if consider_owner_type:
+            return cls._typed_owner_key
+        else:
+            return cls._simple_owner_key
 
-ContainerInputType: TypeAlias = Union["Container", ContainerKey]
+    @classmethod
+    def _typed_owner_key(cls, owner: models.OwnerClass) -> Tuple[str, str]:
+        return (owner.owner, owner.typeUrn or str(owner.type))
+
+    @classmethod
+    def _simple_owner_key(cls, owner: models.OwnerClass) -> Tuple[str,]:
+        return (owner.owner,)
+
+    def _ensure_owners(self) -> List[models.OwnerClass]:
+        owners = self._setdefault_aspect(models.OwnershipClass(owners=[])).owners
+        return owners
+
+    def add_owner(self, owner: OwnerInputType) -> None:
+        # Tricky: when adding an owner, we always use the ownership type.
+        # For removals, we only use it if it was explicitly specified.
+        parsed_owner, _ = self._parse_owner_class(owner)
+        add_list_unique(
+            self._ensure_owners(),
+            key=self._typed_owner_key,
+            item=parsed_owner,
+        )
+
+    def remove_owner(self, owner: OwnerInputType) -> None:
+        parsed_owner, was_type_specified = self._parse_owner_class(owner)
+        remove_list_unique(
+            self._ensure_owners(),
+            key=self._owner_key_method(was_type_specified),
+            item=parsed_owner,
+        )
+
+
+# If you pass in a container object, we can build on top of its browse path.
+# 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]]
 
 
 class HasContainer(Entity):
     __slots__ = ()
 
-    def _set_container(self, container: Optional[ContainerInputType]) -> None:
+    @staticmethod
+    def _maybe_parse_as_urn(urn: UrnOrStr) -> UrnOrStr:
+        if isinstance(urn, Urn):
+            return urn
+        elif urn.startswith("urn:li:"):
+            return Urn.from_string(urn)
+        else:
+            return urn
+
+    def _set_container(self, container: Optional[ParentContainerInputType]) -> 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
 
+        container_urn: Optional[str]
         browse_path: List[Union[str, models.BrowsePathEntryClass]] = []
         if isinstance(container, Container):
             container_urn = container.urn.urn()
@@ -204,6 +269,29 @@ class HasContainer(Entity):
                     urn=container_urn,
                 ),
             ]
+        elif isinstance(container, list):
+            parsed_path = [self._maybe_parse_as_urn(entry) for entry in container]
+
+            # Use the last container in the path as the container urn.
+            container_urns = [
+                urn.urn() for urn in parsed_path if isinstance(urn, ContainerUrn)
+            ]
+            container_urn = container_urns[-1] if container_urns else None
+
+            browse_path = [
+                (
+                    models.BrowsePathEntryClass(
+                        id=str(entry),
+                        urn=str(entry),
+                    )
+                    if isinstance(entry, Urn)
+                    else models.BrowsePathEntryClass(
+                        id=entry,
+                        urn=None,
+                    )
+                )
+                for entry in parsed_path
+            ]
         elif container is not None:
             container_urn = container.as_urn()
 
@@ -243,6 +331,24 @@ class HasContainer(Entity):
             )
         )
 
+    @property
+    def parent_container(self) -> Optional[ContainerUrn]:
+        if container := self._get_aspect(models.ContainerClass):
+            return ContainerUrn.from_string(container.container)
+        return None
+
+    @property
+    def browse_path(self) -> Optional[List[UrnOrStr]]:
+        if browse_path := self._get_aspect(models.BrowsePathsV2Class):
+            path: List[UrnOrStr] = []
+            for entry in browse_path.path:
+                if entry.urn:
+                    path.append(Urn.from_string(entry.urn))
+                else:
+                    path.append(entry.id)
+            return path
+        return None
+
 
 TagInputType: TypeAlias = Union[str, TagUrn, models.TagAssociationClass]
 TagsInputType: TypeAlias = List[TagInputType]
@@ -251,6 +357,9 @@ TagsInputType: TypeAlias = List[TagInputType]
 class HasTags(Entity):
     __slots__ = ()
 
+    def _ensure_tags(self) -> List[models.TagAssociationClass]:
+        return self._setdefault_aspect(models.GlobalTagsClass(tags=[])).tags
+
     # TODO: Return a custom type with deserialized urns, instead of the raw aspect.
     @property
     def tags(self) -> Optional[List[models.TagAssociationClass]]:
@@ -275,6 +384,24 @@ class HasTags(Entity):
             )
         )
 
+    @classmethod
+    def _tag_key(cls, tag: models.TagAssociationClass) -> str:
+        return tag.tag
+
+    def add_tag(self, tag: TagInputType) -> None:
+        add_list_unique(
+            self._ensure_tags(),
+            self._tag_key,
+            self._parse_tag_association_class(tag),
+        )
+
+    def remove_tag(self, tag: TagInputType) -> None:
+        remove_list_unique(
+            self._ensure_tags(),
+            self._tag_key,
+            self._parse_tag_association_class(tag),
+        )
+
 
 TermInputType: TypeAlias = Union[
     str, GlossaryTermUrn, models.GlossaryTermAssociationClass
@@ -285,6 +412,11 @@ TermsInputType: TypeAlias = List[TermInputType]
 class HasTerms(Entity):
     __slots__ = ()
 
+    def _ensure_terms(self) -> List[models.GlossaryTermAssociationClass]:
+        return self._setdefault_aspect(
+            models.GlossaryTermsClass(terms=[], auditStamp=self._terms_audit_stamp())
+        ).terms
+
     # TODO: Return a custom type with deserialized urns, instead of the raw aspect.
     @property
     def terms(self) -> Optional[List[models.GlossaryTermAssociationClass]]:
@@ -320,6 +452,24 @@ class HasTerms(Entity):
             )
         )
 
+    @classmethod
+    def _terms_key(self, term: models.GlossaryTermAssociationClass) -> str:
+        return term.urn
+
+    def add_term(self, term: TermInputType) -> None:
+        add_list_unique(
+            self._ensure_terms(),
+            self._terms_key,
+            self._parse_glossary_term_association_class(term),
+        )
+
+    def remove_term(self, term: TermInputType) -> None:
+        remove_list_unique(
+            self._ensure_terms(),
+            self._terms_key,
+            self._parse_glossary_term_association_class(term),
+        )
+
 
 DomainInputType: TypeAlias = Union[str, DomainUrn]
 

datahub/sdk/_utils.py

@@ -0,0 +1,35 @@
+from typing import Any, Callable, List, Protocol, TypeVar
+
+from datahub.errors import ItemNotFoundError
+
+
+class _SupportsEq(Protocol):
+    def __eq__(self, other: Any) -> bool: ...
+
+
+T = TypeVar("T")
+K = TypeVar("K", bound=_SupportsEq)
+
+
+def add_list_unique(lst: List[T], key: Callable[[T], K], item: T) -> None:
+    item_key = key(item)
+    for i, existing in enumerate(lst):
+        if key(existing) == item_key:
+            lst[i] = item
+            return
+    lst.append(item)
+
+
+def remove_list_unique(
+    lst: List[T], key: Callable[[T], K], item: T, *, missing_ok: bool = True
+) -> None:
+    # Poor man's patch implementation.
+    item_key = key(item)
+    removed = False
+    for i, existing in enumerate(lst):
+        if key(existing) == item_key:
+            lst.pop(i)
+            removed = True
+            # Tricky: no break. In case there's already duplicates, we want to remove all of them.
+    if not removed and not missing_ok:
+        raise ItemNotFoundError(f"Cannot remove item {item} from list: not found")

datahub/sdk/container.py

@@ -16,7 +16,7 @@ from datahub.metadata.urns import (
     ContainerUrn,
     Urn,
 )
-from datahub.sdk._entity import Entity
+from datahub.sdk._entity import Entity, ExtraAspectsType
 from datahub.sdk._shared import (
     DomainInputType,
     HasContainer,
@@ -27,11 +27,13 @@ from datahub.sdk._shared import (
     HasTags,
     HasTerms,
     OwnersInputType,
+    ParentContainerInputType,
     TagsInputType,
     TermsInputType,
     make_time_stamp,
     parse_time_stamp,
 )
+from datahub.utilities.sentinels import Auto, auto
 
 
 class Container(
@@ -54,7 +56,7 @@ class Container(
         self,
         /,
         # Identity.
-        container_key: ContainerKey | ContainerUrn,
+        container_key: ContainerKey,
         *,
         # Container attributes.
         display_name: str,
@@ -66,17 +68,22 @@ class Container(
         created: Optional[datetime] = None,
         last_modified: Optional[datetime] = None,
         # Standard aspects.
+        parent_container: Auto | ParentContainerInputType | None = auto,
         subtype: Optional[str] = None,
         owners: Optional[OwnersInputType] = None,
         tags: Optional[TagsInputType] = None,
         terms: Optional[TermsInputType] = None,
         domain: Optional[DomainInputType] = None,
+        extra_aspects: ExtraAspectsType = None,
     ):
+        # Hack: while the type annotations say container_key is always a ContainerKey,
+        # we allow ContainerUrn to make the graph-based constructor work.
         if isinstance(container_key, ContainerUrn):
             urn = container_key
         else:
             urn = ContainerUrn.from_string(container_key.as_urn())
         super().__init__(urn)
+        self._set_extra_aspects(extra_aspects)
 
         # This needs to come first to ensure that the display name is registered.
         self._ensure_container_props(name=display_name)
@@ -85,8 +92,6 @@ class Container(
         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(),
@@ -100,6 +105,18 @@ class Container(
             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
+        else:
+            self.set_custom_properties(extra_properties or {})
+
+        if parent_container is auto:
+            if not isinstance(container_key, ContainerKey):
+                raise SdkUsageError(
+                    "Either a container_key or parent_container must be provided"
+                )
+
+            self._set_container(container_key.parent_key())
+        else:
+            self._set_container(parent_container)
 
         if description is not None:
             self.set_description(description)
@@ -126,7 +143,8 @@ class Container(
     @classmethod
     def _new_from_graph(cls, urn: Urn, current_aspects: models.AspectBag) -> Self:
         assert isinstance(urn, ContainerUrn)
-        entity = cls(urn, display_name="__dummy_value__")
+
+        entity = cls(urn, display_name="__dummy_value__", parent_container=None)  # type: ignore[arg-type]
         return entity._init_from_graph(current_aspects)
 
     def _ensure_container_props(

datahub/sdk/dataset.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import warnings
 from datetime import datetime
-from typing import Dict, List, Optional, Tuple, Type, Union
+from typing import Dict, List, Optional, Sequence, Tuple, Type, Union
 
 from typing_extensions import Self, TypeAlias, assert_never
 
@@ -13,13 +13,13 @@ from datahub.errors import (
     IngestionAttributionWarning,
     ItemNotFoundError,
     SchemaFieldKeyError,
+    SdkUsageError,
 )
 from datahub.ingestion.source.sql.sql_types import resolve_sql_type
 from datahub.metadata.urns import DatasetUrn, SchemaFieldUrn, Urn
 from datahub.sdk._attribution import is_ingestion_attribution
-from datahub.sdk._entity import Entity
+from datahub.sdk._entity import Entity, ExtraAspectsType
 from datahub.sdk._shared import (
-    ContainerInputType,
     DatasetUrnOrStr,
     DomainInputType,
     HasContainer,
@@ -30,20 +30,24 @@ from datahub.sdk._shared import (
     HasTags,
     HasTerms,
     OwnersInputType,
+    ParentContainerInputType,
+    TagInputType,
     TagsInputType,
+    TermInputType,
     TermsInputType,
     make_time_stamp,
     parse_time_stamp,
 )
+from datahub.sdk._utils import add_list_unique, remove_list_unique
+from datahub.utilities.sentinels import Unset, unset
 
 SchemaFieldInputType: TypeAlias = Union[
-    str,
     Tuple[str, str],  # (name, type)
     Tuple[str, str, str],  # (name, type, description)
     models.SchemaFieldClass,
 ]
 SchemaFieldsInputType: TypeAlias = Union[
-    List[SchemaFieldInputType],
+    Sequence[SchemaFieldInputType],
     models.SchemaMetadataClass,
 ]
 
@@ -271,6 +275,51 @@ class SchemaField:
                 tags=parsed_tags
             )
 
+    def add_tag(self, tag: TagInputType) -> None:
+        parsed_tag = self._parent._parse_tag_association_class(tag)
+
+        if is_ingestion_attribution():
+            raise SdkUsageError(
+                "Adding field tags in ingestion mode is not yet supported. "
+                "Use set_tags instead."
+            )
+        else:
+            editable_field = self._ensure_editable_schema_field()
+            if editable_field.globalTags is None:
+                editable_field.globalTags = models.GlobalTagsClass(tags=[])
+
+            add_list_unique(
+                editable_field.globalTags.tags,
+                key=self._parent._tag_key,
+                item=parsed_tag,
+            )
+
+    def remove_tag(self, tag: TagInputType) -> None:
+        parsed_tag = self._parent._parse_tag_association_class(tag)
+
+        if is_ingestion_attribution():
+            raise SdkUsageError(
+                "Adding field tags in ingestion mode is not yet supported. "
+                "Use set_tags instead."
+            )
+        else:
+            base_field = self._base_schema_field()
+            if base_field.globalTags is not None:
+                remove_list_unique(
+                    base_field.globalTags.tags,
+                    key=self._parent._tag_key,
+                    item=parsed_tag,
+                    missing_ok=True,
+                )
+
+            editable_field = self._ensure_editable_schema_field()
+            if editable_field.globalTags is not None:
+                remove_list_unique(
+                    editable_field.globalTags.tags,
+                    key=self._parent._tag_key,
+                    item=parsed_tag,
+                )
+
     @property
     def terms(self) -> Optional[List[models.GlossaryTermAssociationClass]]:
         # TODO: Basically the same implementation as tags - can we share code?
@@ -287,7 +336,7 @@ class SchemaField:
 
         return terms
 
-    def set_terms(self, terms: List[models.GlossaryTermAssociationClass]) -> None:
+    def set_terms(self, terms: TermsInputType) -> None:
         parsed_terms = [
             self._parent._parse_glossary_term_association_class(term) for term in terms
         ]
@@ -318,6 +367,55 @@ class SchemaField:
                 )
             )
 
+    def add_term(self, term: TermInputType) -> None:
+        parsed_term = self._parent._parse_glossary_term_association_class(term)
+
+        if is_ingestion_attribution():
+            raise SdkUsageError(
+                "Adding field terms in ingestion mode is not yet supported. "
+                "Use set_terms instead."
+            )
+        else:
+            editable_field = self._ensure_editable_schema_field()
+            if editable_field.glossaryTerms is None:
+                editable_field.glossaryTerms = models.GlossaryTermsClass(
+                    terms=[],
+                    auditStamp=self._parent._terms_audit_stamp(),
+                )
+
+            add_list_unique(
+                editable_field.glossaryTerms.terms,
+                key=self._parent._terms_key,
+                item=parsed_term,
+            )
+
+    def remove_term(self, term: TermInputType) -> None:
+        parsed_term = self._parent._parse_glossary_term_association_class(term)
+
+        if is_ingestion_attribution():
+            raise SdkUsageError(
+                "Removing field terms in ingestion mode is not yet supported. "
+                "Use set_terms instead."
+            )
+        else:
+            base_field = self._base_schema_field()
+            if base_field.glossaryTerms is not None:
+                remove_list_unique(
+                    base_field.glossaryTerms.terms,
+                    key=self._parent._terms_key,
+                    item=parsed_term,
+                    missing_ok=True,
+                )
+
+            editable_field = self._ensure_editable_schema_field()
+            if editable_field.glossaryTerms is not None:
+                remove_list_unique(
+                    editable_field.glossaryTerms.terms,
+                    key=self._parent._terms_key,
+                    item=parsed_term,
+                    missing_ok=True,
+                )
+
 
 class Dataset(
     HasPlatformInstance,
@@ -352,13 +450,14 @@ class Dataset(
         created: Optional[datetime] = None,
         last_modified: Optional[datetime] = None,
         # Standard aspects.
+        parent_container: ParentContainerInputType | Unset = unset,
         subtype: Optional[str] = None,
-        container: Optional[ContainerInputType] = None,
         owners: Optional[OwnersInputType] = 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,
@@ -370,6 +469,7 @@ class Dataset(
             env=env,
         )
         super().__init__(urn)
+        self._set_extra_aspects(extra_aspects)
 
         self._set_platform_instance(urn.platform, platform_instance)
 
@@ -393,10 +493,10 @@ class Dataset(
         if last_modified is not None:
             self.set_last_modified(last_modified)
 
+        if parent_container is not unset:
+            self._set_container(parent_container)
         if subtype is not None:
             self.set_subtype(subtype)
-        if container is not None:
-            self._set_container(container)
         if owners is not None:
             self.set_owners(owners)
         if tags is not None:
@@ -537,14 +637,6 @@ class Dataset(
                 nativeDataType=field_type,
                 description=description,
             )
-        elif isinstance(schema_field_input, str):
-            # TODO: Not sure this branch makes sense - we should probably just require types?
-            return models.SchemaFieldClass(
-                fieldPath=schema_field_input,
-                type=models.SchemaFieldDataTypeClass(models.NullTypeClass()),
-                nativeDataType="unknown",
-                description=None,
-            )
         else:
             assert_never(schema_field_input)
 

datahub/sdk/main_client.py

@@ -41,10 +41,24 @@ class DataHubClient:
 
     @classmethod
     def from_env(cls) -> "DataHubClient":
+        """Initialize a DataHubClient from the environment variables or ~/.datahubenv file.
+
+        This will first check DATAHUB_GMS_URL and DATAHUB_GMS_TOKEN. If not present,
+        it will read credentials from ~/.datahubenv. That file can be created using
+        the `datahub init` command.
+
+        If you're looking to specify the server/token in code, use the
+        DataHubClient(server=..., token=...) constructor instead.
+
+        Returns:
+            A DataHubClient instance.
+        """
+
         # Inspired by the DockerClient.from_env() method.
         # TODO: This one also reads from ~/.datahubenv, so the "from_env" name might be a bit confusing.
         # That file is part of the "environment", but is not a traditional "env variable".
         graph = get_default_graph()
+
         return cls(graph=graph)
 
     @property
@@ -54,3 +68,6 @@ class DataHubClient:
     @property
     def resolve(self) -> ResolverClient:
         return ResolverClient(self)
+
+    # TODO: search client
+    # TODO: lineage client