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

datahub/ingestion/source/snowflake/snowflake_utils.py

@@ -124,19 +124,20 @@ class SnowflakeFilter:
             SnowflakeObjectDomain.VIEW,
             SnowflakeObjectDomain.MATERIALIZED_VIEW,
             SnowflakeObjectDomain.ICEBERG_TABLE,
+            SnowflakeObjectDomain.STREAM,
         ):
             return False
         if _is_sys_table(dataset_name):
             return False
 
-        dataset_params = _split_qualified_name(dataset_name)
+        dataset_params = split_qualified_name(dataset_name)
         if len(dataset_params) != 3:
             self.structured_reporter.info(
                 title="Unexpected dataset pattern",
                 message=f"Found a {dataset_type} with an unexpected number of parts. Database and schema filtering will not work as expected, but table filtering will still work.",
                 context=dataset_name,
             )
-            # We fall-through here so table/view filtering still works.
+            # We fall-through here so table/view/stream filtering still works.
 
         if (
             len(dataset_params) >= 1
@@ -169,6 +170,14 @@ class SnowflakeFilter:
         ):
             return False
 
+        if (
+            dataset_type.lower() == SnowflakeObjectDomain.STREAM
+            and not self.filter_config.stream_pattern.allowed(
+                _cleanup_qualified_name(dataset_name, self.structured_reporter)
+            )
+        ):
+            return False
+
         return True
 
 
@@ -183,17 +192,17 @@ def _is_sys_table(table_name: str) -> bool:
     return table_name.lower().startswith("sys$")
 
 
-def _split_qualified_name(qualified_name: str) -> List[str]:
+def split_qualified_name(qualified_name: str) -> List[str]:
     """
     Split a qualified name into its constituent parts.
 
-    >>> _split_qualified_name("db.my_schema.my_table")
+    >>> split_qualified_name("db.my_schema.my_table")
     ['db', 'my_schema', 'my_table']
-    >>> _split_qualified_name('"db"."my_schema"."my_table"')
+    >>> split_qualified_name('"db"."my_schema"."my_table"')
     ['db', 'my_schema', 'my_table']
-    >>> _split_qualified_name('TEST_DB.TEST_SCHEMA."TABLE.WITH.DOTS"')
+    >>> split_qualified_name('TEST_DB.TEST_SCHEMA."TABLE.WITH.DOTS"')
     ['TEST_DB', 'TEST_SCHEMA', 'TABLE.WITH.DOTS']
-    >>> _split_qualified_name('TEST_DB."SCHEMA.WITH.DOTS".MY_TABLE')
+    >>> split_qualified_name('TEST_DB."SCHEMA.WITH.DOTS".MY_TABLE')
     ['TEST_DB', 'SCHEMA.WITH.DOTS', 'MY_TABLE']
     """
 
@@ -231,7 +240,7 @@ def _split_qualified_name(qualified_name: str) -> List[str]:
 def _cleanup_qualified_name(
     qualified_name: str, structured_reporter: SourceReport
 ) -> str:
-    name_parts = _split_qualified_name(qualified_name)
+    name_parts = split_qualified_name(qualified_name)
     if len(name_parts) != 3:
         if not _is_sys_table(qualified_name):
             structured_reporter.info(

datahub/ingestion/source/snowflake/snowflake_v2.py

@@ -539,15 +539,27 @@ class SnowflakeV2Source(
             for schema in db.schemas
             for table_name in schema.views
         ]
+        discovered_streams: List[str] = [
+            self.identifiers.get_dataset_identifier(stream_name, schema.name, db.name)
+            for db in databases
+            for schema in db.schemas
+            for stream_name in schema.streams
+        ]
 
-        if len(discovered_tables) == 0 and len(discovered_views) == 0:
+        if (
+            len(discovered_tables) == 0
+            and len(discovered_views) == 0
+            and len(discovered_streams) == 0
+        ):
             self.structured_reporter.failure(
                 GENERIC_PERMISSION_ERROR_KEY,
-                "No tables/views found. Please check permissions.",
+                "No tables/views/streams found. Please check permissions.",
             )
             return
 
-        self.discovered_datasets = discovered_tables + discovered_views
+        self.discovered_datasets = (
+            discovered_tables + discovered_views + discovered_streams
+        )
 
         if self.config.use_queries_v2:
             with self.report.new_stage(f"*: {VIEW_PARSING}"):

datahub/sdk/__init__.py

@@ -0,0 +1,33 @@
+import warnings
+
+import datahub.metadata.schema_classes as models
+from datahub.errors import ExperimentalWarning, SdkUsageError
+from datahub.ingestion.graph.config import DatahubClientConfig
+from datahub.metadata.urns import (
+    ChartUrn,
+    ContainerUrn,
+    CorpGroupUrn,
+    CorpUserUrn,
+    DashboardUrn,
+    DataPlatformInstanceUrn,
+    DataPlatformUrn,
+    DatasetUrn,
+    DomainUrn,
+    GlossaryTermUrn,
+    SchemaFieldUrn,
+    TagUrn,
+)
+from datahub.sdk.container import Container
+from datahub.sdk.dataset import Dataset
+from datahub.sdk.main_client import DataHubClient
+
+warnings.warn(
+    "The new datahub SDK (e.g. datahub.sdk.*) is experimental. "
+    "Our typical backwards-compatibility and stability guarantees do not apply to this code. "
+    "When it's promoted to stable, the import path will change "
+    "from `from datahub.sdk import ...` to `from datahub import ...`.",
+    ExperimentalWarning,
+    stacklevel=2,
+)
+del warnings
+del ExperimentalWarning

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,338 @@
+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()
+            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)]))