acryl-datahub 0.15.0.6rc3__py3-none-any.whl → 1.0.0__py3-none-any.whl

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,37 +13,43 @@ 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._shared import (
-    ContainerInputType,
     DatasetUrnOrStr,
     DomainInputType,
     HasContainer,
     HasDomain,
+    HasInstitutionalMemory,
     HasOwnership,
     HasPlatformInstance,
     HasSubtype,
     HasTags,
     HasTerms,
+    LinksInputType,
     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.sdk.entity import Entity, ExtraAspectsType
+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,
 ]
 
@@ -68,9 +74,9 @@ UpstreamLineageInputType: TypeAlias = Union[
 def _parse_upstream_input(
     upstream_input: UpstreamInputType,
 ) -> Union[models.UpstreamClass, models.FineGrainedLineageClass]:
-    if isinstance(upstream_input, models.UpstreamClass):
-        return upstream_input
-    elif isinstance(upstream_input, models.FineGrainedLineageClass):
+    if isinstance(
+        upstream_input, (models.UpstreamClass, models.FineGrainedLineageClass)
+    ):
         return upstream_input
     elif isinstance(upstream_input, (str, DatasetUrn)):
         return models.UpstreamClass(
@@ -271,6 +277,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 +338,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,12 +369,62 @@ 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,
     HasSubtype,
     HasContainer,
     HasOwnership,
+    HasInstitutionalMemory,
     HasTags,
     HasTerms,
     HasDomain,
@@ -352,13 +453,15 @@ 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,
+        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,
@@ -370,6 +473,7 @@ class Dataset(
             env=env,
         )
         super().__init__(urn)
+        self._set_extra_aspects(extra_aspects)
 
         self._set_platform_instance(urn.platform, platform_instance)
 
@@ -393,12 +497,14 @@ 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 links is not None:
+            self.set_links(links)
         if tags is not None:
             self.set_tags(tags)
         if terms is not None:
@@ -537,14 +643,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/{_entity.py → entity.py}

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 import abc
-from typing import List, Optional, Type, Union
+from typing import TYPE_CHECKING, List, Optional, Type, Union
 
 from typing_extensions import Self
 
@@ -10,6 +12,12 @@ from datahub.errors import SdkUsageError
 from datahub.metadata.urns import Urn
 from datahub.utilities.urns._urn_base import _SpecificUrn
 
+if TYPE_CHECKING:
+    from datahub.ingestion.api.workunit import MetadataWorkUnit
+
+
+ExtraAspectsType = Union[None, List[AspectTypeVar]]
+
 
 class Entity:
     __slots__ = ("_urn", "_prev_aspects", "_aspects")
@@ -36,6 +44,8 @@ class Entity:
 
     def _init_from_graph(self, current_aspects: models.AspectBag) -> Self:
         self._prev_aspects = current_aspects
+
+        self._aspects = {}
         aspect: models._Aspect
         for aspect_name, aspect in (current_aspects or {}).items():  # type: ignore
             aspect_copy = type(aspect).from_obj(aspect.to_obj())
@@ -46,6 +56,10 @@ class Entity:
     @abc.abstractmethod
     def get_urn_type(cls) -> Type[_SpecificUrn]: ...
 
+    @classmethod
+    def entity_type_name(cls) -> str:
+        return cls.get_urn_type().ENTITY_TYPE
+
     @property
     def urn(self) -> _SpecificUrn:
         return self._urn
@@ -85,5 +99,14 @@ class Entity:
             )
         return mcps
 
+    def as_workunits(self) -> List[MetadataWorkUnit]:
+        return [mcp.as_workunit() for mcp in self._as_mcps()]
+
+    def _set_extra_aspects(self, extra_aspects: ExtraAspectsType) -> None:
+        # TODO: Add validation to ensure that an "extra aspect" does not conflict
+        # with / get overridden by a standard aspect.
+        for aspect in extra_aspects or []:
+            self._set_aspect(aspect)
+
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}('{self.urn}')"

datahub/sdk/entity_client.py

@@ -14,10 +14,10 @@ from datahub.metadata.urns import (
     Urn,
 )
 from datahub.sdk._all_entities import ENTITY_CLASSES
-from datahub.sdk._entity import Entity
 from datahub.sdk._shared import UrnOrStr
 from datahub.sdk.container import Container
 from datahub.sdk.dataset import Dataset
+from datahub.sdk.entity import Entity
 
 if TYPE_CHECKING:
     from datahub.sdk.main_client import DataHubClient

datahub/sdk/main_client.py

@@ -7,6 +7,7 @@ from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
 from datahub.ingestion.graph.config import DatahubClientConfig
 from datahub.sdk.entity_client import EntityClient
 from datahub.sdk.resolver_client import ResolverClient
+from datahub.sdk.search_client import SearchClient
 
 
 class DataHubClient:
@@ -39,12 +40,28 @@ class DataHubClient:
 
         self._graph = graph
 
+    # TODO: test connection
+
     @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 +71,9 @@ class DataHubClient:
     @property
     def resolve(self) -> ResolverClient:
         return ResolverClient(self)
+
+    @property
+    def search(self) -> SearchClient:
+        return SearchClient(self)
+
+    # TODO: lineage client

datahub/sdk/resolver_client.py

@@ -9,6 +9,7 @@ from datahub.metadata.urns import (
     DomainUrn,
     GlossaryTermUrn,
 )
+from datahub.sdk.search_filters import Filter, FilterDsl as F
 
 if TYPE_CHECKING:
     from datahub.sdk.main_client import DataHubClient
@@ -38,37 +39,28 @@ class ResolverClient:
         self, *, name: Optional[str] = None, email: Optional[str] = None
     ) -> CorpUserUrn:
         filter_explanation: str
-        filters = []
+        filter: Filter
         if name is not None:
             if email is not None:
                 raise SdkUsageError("Cannot specify both name and email for auto_user")
-            # TODO: do we filter on displayName or fullName?
+            # We're filtering on both fullName and displayName. It's not clear
+            # what the right behavior is here.
             filter_explanation = f"with name {name}"
-            filters.append(
-                {
-                    "field": "fullName",
-                    "values": [name],
-                    "condition": "EQUAL",
-                }
+            filter = F.or_(
+                F.custom_filter("fullName", "EQUAL", [name]),
+                F.custom_filter("displayName", "EQUAL", [name]),
             )
         elif email is not None:
             filter_explanation = f"with email {email}"
-            filters.append(
-                {
-                    "field": "email",
-                    "values": [email],
-                    "condition": "EQUAL",
-                }
-            )
+            filter = F.custom_filter("email", "EQUAL", [email])
         else:
             raise SdkUsageError("Must specify either name or email for auto_user")
 
-        users = list(
-            self._graph.get_urns_by_filter(
-                entity_types=[CorpUserUrn.ENTITY_TYPE],
-                extraFilters=filters,
-            )
+        filter = F.and_(
+            F.entity_type(CorpUserUrn.ENTITY_TYPE),
+            filter,
         )
+        users = list(self._client.search.get_urns(filter=filter))
         if len(users) == 0:
             # TODO: In auto methods, should we just create the user/domain/etc if it doesn't exist?
             raise ItemNotFoundError(f"User {filter_explanation} not found")
@@ -82,15 +74,11 @@ class ResolverClient:
     def term(self, *, name: str) -> GlossaryTermUrn:
         # TODO: Add some limits on the graph fetch
         terms = list(
-            self._graph.get_urns_by_filter(
-                entity_types=[GlossaryTermUrn.ENTITY_TYPE],
-                extraFilters=[
-                    {
-                        "field": "id",
-                        "values": [name],
-                        "condition": "EQUAL",
-                    }
-                ],
+            self._client.search.get_urns(
+                filter=F.and_(
+                    F.entity_type(GlossaryTermUrn.ENTITY_TYPE),
+                    F.custom_filter("name", "EQUAL", [name]),
+                ),
             )
         )
         if len(terms) == 0:

datahub/sdk/search_client.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+)
+
+from datahub.ingestion.graph.filters import RawSearchFilterRule
+from datahub.metadata.urns import Urn
+from datahub.sdk.search_filters import Filter
+
+if TYPE_CHECKING:
+    from datahub.sdk.main_client import DataHubClient
+
+
+def compile_filters(
+    filter: Optional[Filter],
+) -> Optional[List[Dict[str, List[RawSearchFilterRule]]]]:
+    # TODO: Not every filter type is supported for every entity type.
+    # If we can detect issues with the filters at compile time, we should
+    # raise an error.
+
+    if filter is None:
+        return None
+
+    initial_filters = filter.compile()
+    return [
+        {"and": [rule.to_raw() for rule in andClause["and"]]}
+        for andClause in initial_filters
+    ]
+
+
+class SearchClient:
+    def __init__(self, client: DataHubClient):
+        self._client = client
+
+    def get_urns(
+        self,
+        query: Optional[str] = None,
+        filter: Optional[Filter] = None,
+    ) -> Iterable[Urn]:
+        # TODO: Add better limit / pagination support.
+        for urn in self._client._graph.get_urns_by_filter(
+            query=query,
+            extra_or_filters=compile_filters(filter),
+        ):
+            yield Urn.from_string(urn)