acryl-datahub 1.0.0.2rc1__py3-none-any.whl → 1.0.0.2rc2__py3-none-any.whl

{acryl_datahub-1.0.0.2rc1.dist-info → acryl_datahub-1.0.0.2rc2.dist-info}/RECORD

@@ -1,7 +1,7 @@
-acryl_datahub-1.0.0.2rc1.dist-info/licenses/LICENSE,sha256=9xNHpsD0uYF5ONzXsKDCuHHB-xbiCrSbueWXqrTNsxk,11365
+acryl_datahub-1.0.0.2rc2.dist-info/licenses/LICENSE,sha256=9xNHpsD0uYF5ONzXsKDCuHHB-xbiCrSbueWXqrTNsxk,11365
 datahub/__init__.py,sha256=aq_i5lVREmoLfYIqcx_pEQicO855YlhD19tWc1eZZNI,59
 datahub/__main__.py,sha256=pegIvQ9hzK7IhqVeUi1MeADSZ2QlP-D3K0OQdEg55RU,106
-datahub/_version.py,sha256=vzyBMegu61oWM-Gce9R3y5zLfMrINPSGDEFO-MHhthA,323
+datahub/_version.py,sha256=lFv-ImaIXKL_EDY2GlHJHg9iVkj13C_xihZRNnxH3M8,323
 datahub/entrypoints.py,sha256=2TYgHhs3sCxJlojIHjqfxzt3_ImPwPzq4vBtsUuMqu4,8885
 datahub/errors.py,sha256=BzKdcmYseHOt36zfjJXc17WNutFhp9Y23cU_L6cIkxc,612
 datahub/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -217,7 +217,7 @@ datahub/ingestion/source/redash.py,sha256=YxjSad-X_wPmxYH8dJmFz_VCFhiLTCTSlK99Wd
 datahub/ingestion/source/salesforce.py,sha256=CQtDFv1OsbC1vyzNbKOc6GxhFQ5GdYj45hgAF0-oIcw,40487
 datahub/ingestion/source/source_registry.py,sha256=a2mLjJPLkSI-gYCTb_7U7Jo4D8jGknNQ_yScPIihXFk,1208
 datahub/ingestion/source/sql_queries.py,sha256=Ip7UZub7fgMh7P5jL_zJPY7lSkc9GGTy8GJ8lqZrcsE,9502
-datahub/ingestion/source/superset.py,sha256=FRZ7cCURW6NHUOKaFicdAZq2caXektvO9rJE4tO9scU,40336
+datahub/ingestion/source/superset.py,sha256=bMfvm9HgUoS3T7BjHsDrrOodc8iBRrJRQYv2D66bABo,41194
 datahub/ingestion/source/abs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/abs/config.py,sha256=mBQe0JTaP-Rcv4HnMUUySoYbSr4r3jDEMioxaXHnxXU,6709
 datahub/ingestion/source/abs/datalake_profiler_config.py,sha256=Rkf64evufyVGPiE4VK8QAjzBiJFu85tOGMmJ0lJZ2Og,3600
@@ -328,10 +328,11 @@ datahub/ingestion/source/grafana/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRk
 datahub/ingestion/source/grafana/grafana_source.py,sha256=3pU3xodPgS5lmnjuQ_u7F0XPzD_Y8MnPlMxRJ86qz4g,4960
 datahub/ingestion/source/hex/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/hex/api.py,sha256=JfFPD8O4z16fwZE_BdX5aCQztEq-tbzxJJ7aofH4DE4,12274
-datahub/ingestion/source/hex/constants.py,sha256=NuBjxgJpIt598Cyn_9IcZ158PqBdn5vNjw8T92sTQck,115
-datahub/ingestion/source/hex/hex.py,sha256=DPpsi5e-sdUgbS0Okyvx1mvc00Adu47zA65oFnRP74A,6510
-datahub/ingestion/source/hex/mapper.py,sha256=6dsGvvhPAOAbAG1ayxLwipgJGt1q7YanWYfMX3rZeiM,12603
-datahub/ingestion/source/hex/model.py,sha256=hmMfOLEGZcKjwy2DW29OPf_9_Q_TesgnUTCen2br_fA,1471
+datahub/ingestion/source/hex/constants.py,sha256=8hUTMWyG5keTNfXoLu_Dh413Hw_mGGJX1atiiDZyKtg,271
+datahub/ingestion/source/hex/hex.py,sha256=PIRl8fPkKtlHV7cqR4H8RKVYdTLgEFXHFzc3QAqJLhE,12733
+datahub/ingestion/source/hex/mapper.py,sha256=N3mTlEcrOmhv9ia1dnHGFgFJD2ddyTtU3H5IUbb-UxU,13344
+datahub/ingestion/source/hex/model.py,sha256=S9bUhfFcjzuio2dBS6HzSyRVPiSJvRvMQ0qyVrjV5-E,1766
+datahub/ingestion/source/hex/query_fetcher.py,sha256=5r065vL7XohcgZ_fj-1h6o8cxrPin37IeYsC99GU6LA,12287
 datahub/ingestion/source/iceberg/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/iceberg/iceberg.py,sha256=PhLLXWgBdfZ3hL7LgLvDr6aTK-QKmiZCFNz5jD-mxZM,30773
 datahub/ingestion/source/iceberg/iceberg_common.py,sha256=VGosqYPmn_j6GETSnDHZ8Ay1BVOedmx2x5LHxw16I3A,12278
@@ -1043,8 +1044,8 @@ datahub_provider/operators/datahub_assertion_operator.py,sha256=uvTQ-jk2F0sbqqxp
 datahub_provider/operators/datahub_assertion_sensor.py,sha256=lCBj_3x1cf5GMNpHdfkpHuyHfVxsm6ff5x2Z5iizcAo,140
 datahub_provider/operators/datahub_operation_operator.py,sha256=aevDp2FzX7FxGlXrR0khoHNbxbhKR2qPEX5e8O2Jyzw,174
 datahub_provider/operators/datahub_operation_sensor.py,sha256=8fcdVBCEPgqy1etTXgLoiHoJrRt_nzFZQMdSzHqSG7M,168
-acryl_datahub-1.0.0.2rc1.dist-info/METADATA,sha256=IE26ZK9HREmhmiMf2zQds-JatSIyAh9gcaVjGyOAGLE,176849
-acryl_datahub-1.0.0.2rc1.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-acryl_datahub-1.0.0.2rc1.dist-info/entry_points.txt,sha256=o3mDeJXSKhsy7XLkuogihraiabBdLn9HaizYXPrxmk0,9710
-acryl_datahub-1.0.0.2rc1.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
-acryl_datahub-1.0.0.2rc1.dist-info/RECORD,,
+acryl_datahub-1.0.0.2rc2.dist-info/METADATA,sha256=VuKbVh0Lt8z7Jik8lZ39CF56PZHqn_oIwn2LBmYzrVc,176849
+acryl_datahub-1.0.0.2rc2.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+acryl_datahub-1.0.0.2rc2.dist-info/entry_points.txt,sha256=o3mDeJXSKhsy7XLkuogihraiabBdLn9HaizYXPrxmk0,9710
+acryl_datahub-1.0.0.2rc2.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
+acryl_datahub-1.0.0.2rc2.dist-info/RECORD,,

datahub/_version.py

@@ -1,6 +1,6 @@
 # Published at https://pypi.org/project/acryl-datahub/.
 __package_name__ = "acryl-datahub"
-__version__ = "1.0.0.2rc1"
+__version__ = "1.0.0.2rc2"
 
 
 def is_dev_mode() -> bool:

datahub/ingestion/source/hex/constants.py

@@ -1,3 +1,8 @@
+from datahub.metadata.urns import DataPlatformUrn
+
 HEX_PLATFORM_NAME = "hex"
+HEX_PLATFORM_URN = DataPlatformUrn(platform_name=HEX_PLATFORM_NAME)
 HEX_API_BASE_URL_DEFAULT = "https://app.hex.tech/api/v1"
 HEX_API_PAGE_SIZE_DEFAULT = 100
+
+DATAHUB_API_PAGE_SIZE_DEFAULT = 100

datahub/ingestion/source/hex/hex.py

@@ -1,9 +1,12 @@
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
 from typing import Any, Dict, Iterable, List, Optional
 
-from pydantic import Field, SecretStr
+from pydantic import Field, SecretStr, root_validator
 from typing_extensions import assert_never
 
 from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.datetimes import parse_user_datetime
 from datahub.configuration.source_common import (
     EnvConfigMixin,
     PlatformInstanceConfigMixin,
@@ -21,12 +24,17 @@ from datahub.ingestion.api.source import MetadataWorkUnitProcessor
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.hex.api import HexApi, HexApiReport
 from datahub.ingestion.source.hex.constants import (
+    DATAHUB_API_PAGE_SIZE_DEFAULT,
     HEX_API_BASE_URL_DEFAULT,
     HEX_API_PAGE_SIZE_DEFAULT,
     HEX_PLATFORM_NAME,
 )
 from datahub.ingestion.source.hex.mapper import Mapper
 from datahub.ingestion.source.hex.model import Component, Project
+from datahub.ingestion.source.hex.query_fetcher import (
+    HexQueryFetcher,
+    HexQueryFetcherReport,
+)
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
     StaleEntityRemovalSourceReport,
@@ -34,9 +42,10 @@ from datahub.ingestion.source.state.stale_entity_removal_handler import (
 )
 from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionConfigBase,
-    StatefulIngestionReport,
     StatefulIngestionSourceBase,
 )
+from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
+from datahub.sdk.main_client import DataHubClient
 
 
 class HexSourceConfig(
@@ -93,9 +102,73 @@ class HexSourceConfig(
         default=True,
         description="Set ownership identity from owner/creator email",
     )
+    include_lineage: bool = Field(
+        default=True,
+        description='Include Hex lineage, being fetched from DataHub. See "Limitations" section in the docs for more details about the limitations of this feature.',
+    )
+    lineage_start_time: Optional[datetime] = Field(
+        default=None,
+        description="Earliest date of lineage to consider. Default: 1 day before lineage end time. You can specify absolute time like '2023-01-01' or relative time like '-7 days' or '-7d'.",
+    )
+    lineage_end_time: Optional[datetime] = Field(
+        default=None,
+        description="Latest date of lineage to consider. Default: Current time in UTC. You can specify absolute time like '2023-01-01' or relative time like '-1 day' or '-1d'.",
+    )
+    datahub_page_size: int = Field(
+        default=DATAHUB_API_PAGE_SIZE_DEFAULT,
+        description="Number of items to fetch per DataHub API call.",
+    )
+
+    @root_validator(pre=True)
+    def validate_lineage_times(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        # lineage_end_time default = now
+        if "lineage_end_time" not in data or data["lineage_end_time"] is None:
+            data["lineage_end_time"] = datetime.now(tz=timezone.utc)
+        # if string is given, parse it
+        if isinstance(data["lineage_end_time"], str):
+            data["lineage_end_time"] = parse_user_datetime(data["lineage_end_time"])
+        # if no timezone is given, assume UTC
+        if data["lineage_end_time"].tzinfo is None:
+            data["lineage_end_time"] = data["lineage_end_time"].replace(
+                tzinfo=timezone.utc
+            )
+        # at this point, we ensure there is a non null datetime with UTC timezone for lineage_end_time
+        assert (
+            data["lineage_end_time"]
+            and isinstance(data["lineage_end_time"], datetime)
+            and data["lineage_end_time"].tzinfo is not None
+            and data["lineage_end_time"].tzinfo == timezone.utc
+        )
+
+        # lineage_start_time default = lineage_end_time - 1 day
+        if "lineage_start_time" not in data or data["lineage_start_time"] is None:
+            data["lineage_start_time"] = data["lineage_end_time"] - timedelta(days=1)
+        # if string is given, parse it
+        if isinstance(data["lineage_start_time"], str):
+            data["lineage_start_time"] = parse_user_datetime(data["lineage_start_time"])
+        # if no timezone is given, assume UTC
+        if data["lineage_start_time"].tzinfo is None:
+            data["lineage_start_time"] = data["lineage_start_time"].replace(
+                tzinfo=timezone.utc
+            )
+        # at this point, we ensure there is a non null datetime with UTC timezone for lineage_start_time
+        assert (
+            data["lineage_start_time"]
+            and isinstance(data["lineage_start_time"], datetime)
+            and data["lineage_start_time"].tzinfo is not None
+            and data["lineage_start_time"].tzinfo == timezone.utc
+        )
+
+        return data
 
 
-class HexReport(StaleEntityRemovalSourceReport, HexApiReport):
+@dataclass
+class HexReport(
+    StaleEntityRemovalSourceReport,
+    HexApiReport,
+    IngestionStageReport,
+    HexQueryFetcherReport,
+):
     pass
 
 
@@ -110,7 +183,7 @@ class HexSource(StatefulIngestionSourceBase):
     def __init__(self, config: HexSourceConfig, ctx: PipelineContext):
         super().__init__(config, ctx)
         self.source_config = config
-        self.report = HexReport()
+        self.report: HexReport = HexReport()
         self.platform = HEX_PLATFORM_NAME
         self.hex_api = HexApi(
             report=self.report,
@@ -129,6 +202,28 @@ class HexSource(StatefulIngestionSourceBase):
             categories_as_tags=self.source_config.categories_as_tags,
             set_ownership_from_email=self.source_config.set_ownership_from_email,
         )
+        self.project_registry: Dict[str, Project] = {}
+        self.component_registry: Dict[str, Component] = {}
+
+        self.datahub_client: Optional[DataHubClient] = None
+        self.query_fetcher: Optional[HexQueryFetcher] = None
+        if self.source_config.include_lineage:
+            graph = ctx.require_graph("Lineage")
+            assert self.source_config.lineage_start_time and isinstance(
+                self.source_config.lineage_start_time, datetime
+            )
+            assert self.source_config.lineage_end_time and isinstance(
+                self.source_config.lineage_end_time, datetime
+            )
+            self.datahub_client = DataHubClient(graph=graph)
+            self.query_fetcher = HexQueryFetcher(
+                datahub_client=self.datahub_client,
+                workspace_name=self.source_config.workspace_name,
+                start_datetime=self.source_config.lineage_start_time,
+                end_datetime=self.source_config.lineage_end_time,
+                report=self.report,
+                page_size=self.source_config.datahub_page_size,
+            )
 
     @classmethod
     def create(cls, config_dict: Dict[str, Any], ctx: PipelineContext) -> "HexSource":
@@ -143,25 +238,58 @@ class HexSource(StatefulIngestionSourceBase):
             ).workunit_processor,
         ]
 
-    def get_report(self) -> StatefulIngestionReport:
+    def get_report(self) -> HexReport:
         return self.report
 
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
-        yield from self.mapper.map_workspace()
-
-        for project_or_component in self.hex_api.fetch_projects():
-            if isinstance(project_or_component, Project):
-                if self.source_config.project_title_pattern.allowed(
-                    project_or_component.title
-                ):
-                    yield from self.mapper.map_project(project=project_or_component)
-            elif isinstance(project_or_component, Component):
-                if (
-                    self.source_config.include_components
-                    and self.source_config.component_title_pattern.allowed(
+        with self.report.new_stage("Fetch Hex assets from Hex API"):
+            for project_or_component in self.hex_api.fetch_projects():
+                if isinstance(project_or_component, Project):
+                    if self.source_config.project_title_pattern.allowed(
                         project_or_component.title
-                    )
-                ):
-                    yield from self.mapper.map_component(component=project_or_component)
-            else:
-                assert_never(project_or_component)
+                    ):
+                        self.project_registry[project_or_component.id] = (
+                            project_or_component
+                        )
+                elif isinstance(project_or_component, Component):
+                    if (
+                        self.source_config.include_components
+                        and self.source_config.component_title_pattern.allowed(
+                            project_or_component.title
+                        )
+                    ):
+                        self.component_registry[project_or_component.id] = (
+                            project_or_component
+                        )
+                else:
+                    assert_never(project_or_component)
+
+        if self.source_config.include_lineage:
+            assert self.datahub_client and self.query_fetcher
+
+            with self.report.new_stage(
+                "Fetch Hex lineage from existing Queries in DataHub"
+            ):
+                for query_metadata in self.query_fetcher.fetch():
+                    project = self.project_registry.get(query_metadata.hex_project_id)
+                    if project:
+                        project.upstream_datasets.extend(
+                            query_metadata.dataset_subjects
+                        )
+                        project.upstream_schema_fields.extend(
+                            query_metadata.schema_field_subjects
+                        )
+                    else:
+                        self.report.report_warning(
+                            title="Missing project for lineage",
+                            message="Lineage missed because missed project, likely due to filter patterns or deleted project.",
+                            context=str(query_metadata),
+                        )
+
+        with self.report.new_stage("Emit"):
+            yield from self.mapper.map_workspace()
+
+            for project in self.project_registry.values():
+                yield from self.mapper.map_project(project=project)
+            for component in self.component_registry.values():
+                yield from self.mapper.map_component(component=component)

datahub/ingestion/source/hex/mapper.py

@@ -1,6 +1,6 @@
 import logging
 from datetime import datetime
-from typing import Iterable, List, Optional, Tuple
+from typing import Iterable, List, Optional, Tuple, Union
 
 from datahub._codegen.aspect import (
     _Aspect,  # TODO: is there a better import than this one?
@@ -46,6 +46,7 @@ from datahub.metadata.schema_classes import (
     DashboardInfoClass,
     DashboardUsageStatisticsClass,
     DataPlatformInstanceClass,
+    EdgeClass,
     GlobalTagsClass,
     OwnerClass,
     OwnershipClass,
@@ -53,7 +54,14 @@ from datahub.metadata.schema_classes import (
     TagAssociationClass,
     TimeWindowSizeClass,
 )
-from datahub.metadata.urns import ContainerUrn, CorpUserUrn, DashboardUrn, Urn
+from datahub.metadata.urns import (
+    ContainerUrn,
+    CorpUserUrn,
+    DashboardUrn,
+    DatasetUrn,
+    SchemaFieldUrn,
+    Urn,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -116,6 +124,8 @@ class Mapper:
             ),
             externalUrl=f"{self._base_url}/{self._workspace_name}/hex/{project.id}",
             customProperties=dict(id=project.id),
+            datasetEdges=self._dataset_edges(project.upstream_datasets),
+            # TODO: support schema field upstream, maybe InputFields?
         )
 
         subtypes = SubTypesClass(
@@ -343,6 +353,22 @@ class Mapper:
             else None,
         )
 
+    def _dataset_edges(
+        self, upstream: List[Union[DatasetUrn, SchemaFieldUrn]]
+    ) -> Optional[List[EdgeClass]]:
+        # TBC: is there support for CLL in Dashboards? for the moment, skip SchemaFieldUrns
+        return (
+            [
+                EdgeClass(
+                    destinationUrn=upstream_urn.urn(),
+                )
+                for upstream_urn in upstream
+                if isinstance(upstream_urn, DatasetUrn)
+            ]
+            if upstream
+            else None
+        )
+
     def _yield_mcps(
         self, entity_urn: Urn, aspects: List[Optional[_Aspect]]
     ) -> Iterable[MetadataWorkUnit]:

datahub/ingestion/source/hex/model.py

@@ -1,6 +1,8 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import datetime
-from typing import List, Optional
+from typing import List, Optional, Union
+
+from datahub.metadata.urns import DatasetUrn, SchemaFieldUrn
 
 
 @dataclass
@@ -51,6 +53,12 @@ class Project:
     creator: Optional[Owner] = None
     owner: Optional[Owner] = None
     analytics: Optional[Analytics] = None
+    upstream_datasets: List[Union[DatasetUrn, SchemaFieldUrn]] = field(
+        default_factory=list
+    )
+    upstream_schema_fields: List[Union[DatasetUrn, SchemaFieldUrn]] = field(
+        default_factory=list
+    )
 
 
 @dataclass