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

datahub/ingestion/source/hex/mapper.py

@@ -0,0 +1,372 @@
+import logging
+from datetime import datetime
+from typing import Iterable, List, Optional, Tuple
+
+from datahub._codegen.aspect import (
+    _Aspect,  # TODO: is there a better import than this one?
+)
+from datahub.emitter.mce_builder import (
+    make_container_urn,
+    make_dashboard_urn,
+    make_data_platform_urn,
+    make_dataplatform_instance_urn,
+    make_tag_urn,
+    make_ts_millis,
+    make_user_urn,
+)
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.emitter.mcp_builder import ContainerKey
+from datahub.ingestion.api.incremental_lineage_helper import (
+    convert_dashboard_info_to_patch,
+)
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import BIAssetSubTypes
+from datahub.ingestion.source.hex.constants import (
+    HEX_API_BASE_URL_DEFAULT,
+    HEX_PLATFORM_NAME,
+)
+from datahub.ingestion.source.hex.model import (
+    Analytics,
+    Category,
+    Collection,
+    Component,
+    Owner,
+    Project,
+    Status,
+)
+from datahub.metadata.com.linkedin.pegasus2avro.common import (
+    AuditStampClass,
+    ChangeAuditStampsClass,
+    OwnershipType,
+)
+from datahub.metadata.schema_classes import (
+    CalendarIntervalClass,
+    ContainerClass,
+    ContainerPropertiesClass,
+    DashboardInfoClass,
+    DashboardUsageStatisticsClass,
+    DataPlatformInstanceClass,
+    GlobalTagsClass,
+    OwnerClass,
+    OwnershipClass,
+    SubTypesClass,
+    TagAssociationClass,
+    TimeWindowSizeClass,
+)
+from datahub.metadata.urns import ContainerUrn, CorpUserUrn, DashboardUrn, Urn
+
+logger = logging.getLogger(__name__)
+
+
+class WorkspaceKey(ContainerKey):
+    workspace_name: str
+
+
+DEFAULT_INGESTION_USER_URN = CorpUserUrn("_ingestion")
+DEFAULT_OWNERSHIP_TYPE = OwnershipType.TECHNICAL_OWNER
+
+
+class Mapper:
+    def __init__(
+        self,
+        workspace_name: str,
+        platform_instance: Optional[str] = None,
+        env: Optional[str] = None,
+        base_url: str = HEX_API_BASE_URL_DEFAULT,
+        patch_metadata: bool = True,
+        collections_as_tags: bool = True,
+        status_as_tag: bool = True,
+        categories_as_tags: bool = True,
+        set_ownership_from_email: bool = True,
+    ):
+        self._workspace_name = workspace_name
+        self._env = env
+        self._platform_instance = platform_instance
+        self._workspace_urn = Mapper._get_workspace_urn(
+            workspace_name=workspace_name,
+            platform=HEX_PLATFORM_NAME,
+            env=env,
+            platform_instance=platform_instance,
+        )
+        self._base_url = base_url.strip("/").replace("/api/v1", "")
+        self._patch_metadata = patch_metadata
+        self._collections_as_tags = collections_as_tags
+        self._status_as_tag = status_as_tag
+        self._categories_as_tags = categories_as_tags
+        self._set_ownership_from_email = set_ownership_from_email
+
+    def map_workspace(self) -> Iterable[MetadataWorkUnit]:
+        container_properties = ContainerPropertiesClass(
+            name=self._workspace_name,
+            env=self._env,
+        )
+        yield from self._yield_mcps(
+            entity_urn=self._workspace_urn,
+            aspects=[container_properties],
+        )
+
+    def map_project(self, project: Project) -> Iterable[MetadataWorkUnit]:
+        dashboard_urn = self._get_dashboard_urn(name=project.id)
+
+        dashboard_info = DashboardInfoClass(
+            title=project.title,
+            description=project.description or "",
+            lastModified=self._change_audit_stamps(
+                created_at=project.created_at, last_edited_at=project.last_edited_at
+            ),
+            externalUrl=f"{self._base_url}/{self._workspace_name}/hex/{project.id}",
+            customProperties=dict(id=project.id),
+        )
+
+        subtypes = SubTypesClass(
+            typeNames=[BIAssetSubTypes.HEX_PROJECT],
+        )
+
+        platform_instance = self._platform_instance_aspect()
+
+        container = ContainerClass(
+            container=self._workspace_urn.urn(),
+        )
+
+        tags = self._global_tags(
+            status=project.status,
+            categories=project.categories,
+            collections=project.collections,
+        )
+
+        ownership = self._ownership(creator=project.creator, owner=project.owner)
+
+        usage_stats_all_time, usage_stats_last_7_days = (
+            self._dashboard_usage_statistics(analytics=project.analytics)
+        )
+
+        yield from self._yield_mcps(
+            entity_urn=dashboard_urn,
+            aspects=[
+                dashboard_info,
+                subtypes,
+                platform_instance,
+                container,
+                tags,
+                ownership,
+                usage_stats_all_time,
+                usage_stats_last_7_days,
+            ],
+        )
+
+    def map_component(self, component: Component) -> Iterable[MetadataWorkUnit]:
+        dashboard_urn = self._get_dashboard_urn(name=component.id)
+
+        dashboard_info = DashboardInfoClass(
+            title=component.title,
+            description=component.description or "",
+            lastModified=self._change_audit_stamps(
+                created_at=component.created_at, last_edited_at=component.last_edited_at
+            ),
+            externalUrl=f"{self._base_url}/{self._workspace_name}/hex/{component.id}",
+            customProperties=dict(id=component.id),
+        )
+
+        subtypes = SubTypesClass(
+            typeNames=[BIAssetSubTypes.HEX_COMPONENT],
+        )
+
+        platform_instance = self._platform_instance_aspect()
+
+        container = ContainerClass(
+            container=self._workspace_urn.urn(),
+        )
+
+        tags = self._global_tags(
+            status=component.status,
+            categories=component.categories,
+            collections=component.collections,
+        )
+
+        ownership = self._ownership(creator=component.creator, owner=component.owner)
+
+        usage_stats_all_time, usage_stats_last_7_days = (
+            self._dashboard_usage_statistics(analytics=component.analytics)
+        )
+
+        yield from self._yield_mcps(
+            entity_urn=dashboard_urn,
+            aspects=[
+                dashboard_info,
+                subtypes,
+                platform_instance,
+                container,
+                tags,
+                ownership,
+                usage_stats_all_time,
+                usage_stats_last_7_days,
+            ],
+        )
+
+    @classmethod
+    def _get_workspace_urn(
+        cls,
+        workspace_name: str,
+        platform: str = HEX_PLATFORM_NAME,
+        env: Optional[str] = None,
+        platform_instance: Optional[str] = None,
+    ) -> ContainerUrn:
+        workspace_key = WorkspaceKey(
+            platform=platform,
+            env=env,
+            platform_instance=platform_instance,
+            workspace_name=workspace_name,
+        )
+        container_urn_str = make_container_urn(guid=workspace_key)
+        container_urn = Urn.from_string(container_urn_str)
+        assert isinstance(container_urn, ContainerUrn)
+        return container_urn
+
+    def _get_dashboard_urn(self, name: str) -> DashboardUrn:
+        dashboard_urn_str = make_dashboard_urn(
+            platform=HEX_PLATFORM_NAME,
+            name=name,
+            platform_instance=self._platform_instance,
+        )
+        dashboard_urn = Urn.from_string(dashboard_urn_str)
+        assert isinstance(dashboard_urn, DashboardUrn)
+        return dashboard_urn
+
+    def _change_audit_stamps(
+        self, created_at: Optional[datetime], last_edited_at: Optional[datetime]
+    ) -> ChangeAuditStampsClass:
+        return ChangeAuditStampsClass(
+            created=AuditStampClass(
+                time=make_ts_millis(created_at),
+                actor=DEFAULT_INGESTION_USER_URN.urn(),
+            )
+            if created_at
+            else None,
+            lastModified=AuditStampClass(
+                time=make_ts_millis(last_edited_at),
+                actor=DEFAULT_INGESTION_USER_URN.urn(),
+            )
+            if last_edited_at
+            else None,
+        )
+
+    def _global_tags(
+        self,
+        status: Optional[Status],
+        categories: Optional[List[Category]],
+        collections: Optional[List[Collection]],
+    ) -> Optional[GlobalTagsClass]:
+        tag_associations: List[TagAssociationClass] = []
+        if status and self._status_as_tag:
+            tag_associations.append(
+                TagAssociationClass(tag=make_tag_urn(tag=f"hex:status:{status.name}"))
+            )
+        if categories and self._categories_as_tags:
+            tag_associations.extend(
+                [
+                    TagAssociationClass(
+                        tag=make_tag_urn(tag=f"hex:category:{cat.name}")
+                    )
+                    for cat in categories
+                ]
+            )
+        if collections and self._collections_as_tags:
+            tag_associations.extend(
+                [
+                    TagAssociationClass(
+                        tag=make_tag_urn(tag=f"hex:collection:{col.name}")
+                    )
+                    for col in collections
+                ]
+            )
+
+        return GlobalTagsClass(tags=tag_associations) if tag_associations else None
+
+    def _ownership(
+        self, creator: Optional[Owner], owner: Optional[Owner]
+    ) -> Optional[OwnershipClass]:
+        if self._set_ownership_from_email:
+            # since we are not making any diff of creator/owner, we usually have duplicates
+            # TODO: set or ownership types to properly differentiate them, maybe by config?
+            unique_owners = set(o for o in [creator, owner] if o)
+            owners: List[OwnerClass] = [
+                OwnerClass(owner=make_user_urn(o.email), type=DEFAULT_OWNERSHIP_TYPE)
+                for o in unique_owners
+            ]
+
+            return OwnershipClass(owners=owners) if owners else None
+        return None
+
+    def _dashboard_usage_statistics(
+        self, analytics: Optional[Analytics]
+    ) -> Tuple[
+        Optional[DashboardUsageStatisticsClass], Optional[DashboardUsageStatisticsClass]
+    ]:
+        tm_millis = make_ts_millis(datetime.now())
+        last_viewed_at = (
+            make_ts_millis(analytics.last_viewed_at)
+            if analytics and analytics.last_viewed_at
+            else None
+        )
+
+        usage_all_time: Optional[DashboardUsageStatisticsClass] = (
+            DashboardUsageStatisticsClass(
+                timestampMillis=tm_millis,
+                viewsCount=analytics.appviews_all_time,
+                lastViewedAt=last_viewed_at,
+            )
+            if analytics and analytics.appviews_all_time
+            else None
+        )
+
+        usage_last_7_days: Optional[DashboardUsageStatisticsClass] = (
+            DashboardUsageStatisticsClass(
+                timestampMillis=tm_millis,
+                viewsCount=analytics.appviews_last_7_days,
+                eventGranularity=TimeWindowSizeClass(
+                    unit=CalendarIntervalClass.WEEK, multiple=1
+                ),
+                lastViewedAt=last_viewed_at,
+            )
+            if analytics and analytics.appviews_last_7_days
+            else None
+        )
+        return (usage_all_time, usage_last_7_days)
+
+    def _platform_instance_aspect(self) -> DataPlatformInstanceClass:
+        return DataPlatformInstanceClass(
+            platform=make_data_platform_urn(HEX_PLATFORM_NAME),
+            instance=make_dataplatform_instance_urn(
+                platform=HEX_PLATFORM_NAME, instance=self._platform_instance
+            )
+            if self._platform_instance
+            else None,
+        )
+
+    def _yield_mcps(
+        self, entity_urn: Urn, aspects: List[Optional[_Aspect]]
+    ) -> Iterable[MetadataWorkUnit]:
+        for mcpw in MetadataChangeProposalWrapper.construct_many(
+            entityUrn=entity_urn.urn(),
+            aspects=aspects,
+        ):
+            wu = MetadataWorkUnit.from_metadata(metadata=mcpw)
+            maybe_wu = self._maybe_patch_wu(wu)
+            if maybe_wu:
+                yield maybe_wu
+
+    def _maybe_patch_wu(self, wu: MetadataWorkUnit) -> Optional[MetadataWorkUnit]:
+        # So far we only have support for DashboardInfo aspect
+
+        dashboard_info_aspect: Optional[DashboardInfoClass] = wu.get_aspect_of_type(
+            DashboardInfoClass
+        )
+
+        if dashboard_info_aspect and self._patch_metadata:
+            return convert_dashboard_info_to_patch(
+                wu.get_urn(),
+                dashboard_info_aspect,
+                wu.metadata.systemMetadata,
+            )
+        else:
+            return wu

datahub/ingestion/source/hex/model.py

@@ -0,0 +1,68 @@
+from dataclasses import dataclass
+from datetime import datetime
+from typing import List, Optional
+
+
+@dataclass
+class Workspace:
+    name: str
+
+
+@dataclass
+class Status:
+    name: str
+
+
+@dataclass
+class Category:
+    name: str
+    description: Optional[str] = None
+
+
+@dataclass
+class Collection:
+    name: str
+
+
+@dataclass(frozen=True)
+class Owner:
+    email: str
+
+
+@dataclass
+class Analytics:
+    appviews_all_time: Optional[int]
+    appviews_last_7_days: Optional[int]
+    appviews_last_14_days: Optional[int]
+    appviews_last_30_days: Optional[int]
+    last_viewed_at: Optional[datetime]
+
+
+@dataclass
+class Project:
+    id: str
+    title: str
+    description: Optional[str]
+    last_edited_at: Optional[datetime] = None
+    created_at: Optional[datetime] = None
+    status: Optional[Status] = None
+    categories: Optional[List[Category]] = None  # TODO: emit category description!
+    collections: Optional[List[Collection]] = None
+    creator: Optional[Owner] = None
+    owner: Optional[Owner] = None
+    analytics: Optional[Analytics] = None
+
+
+@dataclass
+class Component:
+    id: str
+    title: str
+    description: Optional[str]
+    last_edited_at: Optional[datetime] = None
+    created_at: Optional[datetime] = None
+    status: Optional[Status] = None
+    categories: Optional[List[Category]] = None
+    collections: Optional[List[Collection]] = None
+    creator: Optional[Owner] = None
+    owner: Optional[Owner] = None
+    analytics: Optional[Analytics] = None

datahub/ingestion/source/iceberg/iceberg.py

@@ -69,8 +69,6 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionSourceBase,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.common import Status
-from datahub.metadata.com.linkedin.pegasus2avro.metadata.snapshot import DatasetSnapshot
-from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
 from datahub.metadata.com.linkedin.pegasus2avro.schema import (
     OtherSchema,
     SchemaField,
@@ -298,67 +296,35 @@ class IcebergSource(StatefulIngestionSourceBase):
                 self.config.platform_instance,
                 self.config.env,
             )
-            dataset_snapshot = DatasetSnapshot(
-                urn=dataset_urn,
-                aspects=[Status(removed=False)],
-            )
-
-            # Dataset properties aspect.
-            additional_properties = {}
-            custom_properties = table.metadata.properties.copy()
-            custom_properties["location"] = table.metadata.location
-            custom_properties["format-version"] = str(table.metadata.format_version)
-            custom_properties["partition-spec"] = str(self._get_partition_aspect(table))
-            if table.current_snapshot():
-                custom_properties["snapshot-id"] = str(
-                    table.current_snapshot().snapshot_id
-                )
-                custom_properties["manifest-list"] = (
-                    table.current_snapshot().manifest_list
-                )
-                additional_properties["lastModified"] = TimeStampClass(
-                    int(table.current_snapshot().timestamp_ms)
-                )
-            if "created-at" in custom_properties:
-                try:
-                    dt = dateutil_parser.isoparse(custom_properties["created-at"])
-                    additional_properties["created"] = TimeStampClass(
-                        int(dt.timestamp() * 1000)
-                    )
-                except Exception as ex:
-                    LOGGER.warning(
-                        f"Exception while trying to parse creation date {custom_properties['created-at']}, ignoring: {ex}"
-                    )
+            yield MetadataChangeProposalWrapper(
+                entityUrn=dataset_urn, aspect=Status(removed=False)
+            ).as_workunit()
 
-            dataset_properties = DatasetPropertiesClass(
-                name=table.name()[-1],
-                description=table.metadata.properties.get("comment", None),
-                customProperties=custom_properties,
-                lastModified=additional_properties.get("lastModified"),
-                created=additional_properties.get("created"),
-                qualifiedName=dataset_name,
+            dataset_properties = self._get_dataset_properties_aspect(
+                dataset_name, table
             )
-            dataset_snapshot.aspects.append(dataset_properties)
-            # Dataset ownership aspect.
+            yield MetadataChangeProposalWrapper(
+                entityUrn=dataset_urn, aspect=dataset_properties
+            ).as_workunit()
+
             dataset_ownership = self._get_ownership_aspect(table)
             if dataset_ownership:
                 LOGGER.debug(
                     f"Adding ownership: {dataset_ownership} to the dataset {dataset_name}"
                 )
-                dataset_snapshot.aspects.append(dataset_ownership)
+                yield MetadataChangeProposalWrapper(
+                    entityUrn=dataset_urn, aspect=dataset_ownership
+                ).as_workunit()
 
             schema_metadata = self._create_schema_metadata(dataset_name, table)
-            dataset_snapshot.aspects.append(schema_metadata)
+            yield MetadataChangeProposalWrapper(
+                entityUrn=dataset_urn, aspect=schema_metadata
+            ).as_workunit()
+            yield self._get_dataplatform_instance_aspect(dataset_urn=dataset_urn)
 
-            mce = MetadataChangeEvent(proposedSnapshot=dataset_snapshot)
         self.report.report_table_processing_time(
             timer.elapsed_seconds(), dataset_name, table.metadata_location
         )
-        yield MetadataWorkUnit(id=dataset_name, mce=mce)
-
-        dpi_aspect = self._get_dataplatform_instance_aspect(dataset_urn=dataset_urn)
-        if dpi_aspect:
-            yield dpi_aspect
 
         if self.config.is_profiling_enabled():
             profiler = IcebergProfiler(self.report, self.config.profiling)
@@ -407,6 +373,40 @@ class IcebergSource(StatefulIngestionSourceBase):
             )
             return None
 
+    def _get_dataset_properties_aspect(
+        self, dataset_name: str, table: Table
+    ) -> DatasetPropertiesClass:
+        additional_properties = {}
+        custom_properties = table.metadata.properties.copy()
+        custom_properties["location"] = table.metadata.location
+        custom_properties["format-version"] = str(table.metadata.format_version)
+        custom_properties["partition-spec"] = str(self._get_partition_aspect(table))
+        if table.current_snapshot():
+            custom_properties["snapshot-id"] = str(table.current_snapshot().snapshot_id)
+            custom_properties["manifest-list"] = table.current_snapshot().manifest_list
+            additional_properties["lastModified"] = TimeStampClass(
+                int(table.current_snapshot().timestamp_ms)
+            )
+        if "created-at" in custom_properties:
+            try:
+                dt = dateutil_parser.isoparse(custom_properties["created-at"])
+                additional_properties["created"] = TimeStampClass(
+                    int(dt.timestamp() * 1000)
+                )
+            except Exception as ex:
+                LOGGER.warning(
+                    f"Exception while trying to parse creation date {custom_properties['created-at']}, ignoring: {ex}"
+                )
+
+        return DatasetPropertiesClass(
+            name=table.name()[-1],
+            description=table.metadata.properties.get("comment", None),
+            customProperties=custom_properties,
+            lastModified=additional_properties.get("lastModified"),
+            created=additional_properties.get("created"),
+            qualifiedName=dataset_name,
+        )
+
     def _get_ownership_aspect(self, table: Table) -> Optional[OwnershipClass]:
         owners = []
         if self.config.user_ownership_property:
@@ -435,22 +435,18 @@ class IcebergSource(StatefulIngestionSourceBase):
                 )
         return OwnershipClass(owners=owners) if owners else None
 
-    def _get_dataplatform_instance_aspect(
-        self, dataset_urn: str
-    ) -> Optional[MetadataWorkUnit]:
-        # If we are a platform instance based source, emit the instance aspect
-        if self.config.platform_instance:
-            return MetadataChangeProposalWrapper(
-                entityUrn=dataset_urn,
-                aspect=DataPlatformInstanceClass(
-                    platform=make_data_platform_urn(self.platform),
-                    instance=make_dataplatform_instance_urn(
-                        self.platform, self.config.platform_instance
-                    ),
-                ),
-            ).as_workunit()
-
-        return None
+    def _get_dataplatform_instance_aspect(self, dataset_urn: str) -> MetadataWorkUnit:
+        return MetadataChangeProposalWrapper(
+            entityUrn=dataset_urn,
+            aspect=DataPlatformInstanceClass(
+                platform=make_data_platform_urn(self.platform),
+                instance=make_dataplatform_instance_urn(
+                    self.platform, self.config.platform_instance
+                )
+                if self.config.platform_instance
+                else None,
+            ),
+        ).as_workunit()
 
     def _create_schema_metadata(
         self, dataset_name: str, table: Table