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

datahub/ingestion/source/unity/source.py

@@ -7,12 +7,14 @@ from urllib.parse import urljoin
 from datahub.api.entities.external.external_entities import PlatformResourceRepository
 from datahub.api.entities.external.unity_catalog_external_entites import UnityCatalogTag
 from datahub.emitter.mce_builder import (
+    UNKNOWN_USER,
     make_data_platform_urn,
     make_dataplatform_instance_urn,
     make_dataset_urn_with_platform_instance,
     make_domain_urn,
     make_group_urn,
     make_schema_field_urn,
+    make_ts_millis,
     make_user_urn,
 )
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
@@ -111,6 +113,7 @@ from datahub.metadata.com.linkedin.pegasus2avro.dataset import (
     ViewProperties,
 )
 from datahub.metadata.schema_classes import (
+    AuditStampClass,
     BrowsePathsClass,
     DataPlatformInstanceClass,
     DatasetLineageTypeClass,
@@ -203,6 +206,7 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             config.warehouse_id,
             report=self.report,
             hive_metastore_proxy=self.hive_metastore_proxy,
+            lineage_data_source=config.lineage_data_source,
         )
 
         self.external_url_base = urljoin(self.config.workspace_url, "/explore/data")
@@ -410,12 +414,12 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
                         self.config.workspace_url, f"#notebook/{notebook.id}"
                     ),
                     created=(
-                        TimeStampClass(int(notebook.created_at.timestamp() * 1000))
+                        TimeStampClass(make_ts_millis(notebook.created_at))
                         if notebook.created_at
                         else None
                     ),
                     lastModified=(
-                        TimeStampClass(int(notebook.modified_at.timestamp() * 1000))
+                        TimeStampClass(make_ts_millis(notebook.modified_at))
                         if notebook.modified_at
                         else None
                     ),
@@ -434,17 +438,20 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
         if not notebook.upstreams:
             return None
 
+        upstreams = []
+        for upstream_ref in notebook.upstreams:
+            timestamp = make_ts_millis(upstream_ref.last_updated)
+            upstreams.append(
+                self._create_upstream_class(
+                    self.gen_dataset_urn(upstream_ref),
+                    DatasetLineageTypeClass.COPY,
+                    timestamp,
+                )
+            )
+
         return MetadataChangeProposalWrapper(
             entityUrn=self.gen_notebook_urn(notebook),
-            aspect=UpstreamLineageClass(
-                upstreams=[
-                    UpstreamClass(
-                        dataset=self.gen_dataset_urn(upstream_ref),
-                        type=DatasetLineageTypeClass.COPY,
-                    )
-                    for upstream_ref in notebook.upstreams
-                ]
-            ),
+            aspect=UpstreamLineageClass(upstreams=upstreams),
         ).as_workunit()
 
     def process_metastores(self) -> Iterable[MetadataWorkUnit]:
@@ -463,14 +470,15 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
         self, metastore: Optional[Metastore]
     ) -> Iterable[MetadataWorkUnit]:
         for catalog in self._get_catalogs(metastore):
-            if not self.config.catalog_pattern.allowed(catalog.id):
-                self.report.catalogs.dropped(catalog.id)
-                continue
+            with self.report.new_stage(f"Ingest catalog {catalog.id}"):
+                if not self.config.catalog_pattern.allowed(catalog.id):
+                    self.report.catalogs.dropped(catalog.id)
+                    continue
 
-            yield from self.gen_catalog_containers(catalog)
-            yield from self.process_schemas(catalog)
+                yield from self.gen_catalog_containers(catalog)
+                yield from self.process_schemas(catalog)
 
-            self.report.catalogs.processed(catalog.id)
+                self.report.catalogs.processed(catalog.id)
 
     def _get_catalogs(self, metastore: Optional[Metastore]) -> Iterable[Catalog]:
         if self.config.catalogs:
@@ -647,9 +655,21 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
         ]
 
     def ingest_lineage(self, table: Table) -> Optional[UpstreamLineageClass]:
+        # Calculate datetime filters for lineage
+        lineage_start_time = None
+        lineage_end_time = self.config.end_time
+
+        if self.config.ignore_start_time_lineage:
+            lineage_start_time = None  # Ignore start time to get all lineage
+        else:
+            lineage_start_time = self.config.start_time
+
         if self.config.include_table_lineage:
             self.unity_catalog_api_proxy.table_lineage(
-                table, include_entity_lineage=self.config.include_notebooks
+                table,
+                include_entity_lineage=self.config.include_notebooks,
+                start_time=lineage_start_time,
+                end_time=lineage_end_time,
             )
 
         if self.config.include_column_lineage and table.upstreams:
@@ -661,7 +681,11 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
                 for column in table.columns[: self.config.column_lineage_column_limit]
             ]
             self.unity_catalog_api_proxy.get_column_lineage(
-                table, column_names, max_workers=self.config.lineage_max_workers
+                table,
+                column_names,
+                max_workers=self.config.lineage_max_workers,
+                start_time=lineage_start_time,
+                end_time=lineage_end_time,
             )
 
         return self._generate_lineage_aspect(self.gen_dataset_urn(table.ref), table)
@@ -690,18 +714,22 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
                 for d_col, u_cols in sorted(downstream_to_upstream_cols.items())
             )
 
+            timestamp = make_ts_millis(upstream_ref.last_updated)
             upstreams.append(
-                UpstreamClass(
-                    dataset=upstream_urn,
-                    type=DatasetLineageTypeClass.TRANSFORMED,
+                self._create_upstream_class(
+                    upstream_urn,
+                    DatasetLineageTypeClass.TRANSFORMED,
+                    timestamp,
                 )
             )
 
-        for notebook in table.upstream_notebooks:
+        for notebook in table.upstream_notebooks.values():
+            timestamp = make_ts_millis(notebook.last_updated)
             upstreams.append(
-                UpstreamClass(
-                    dataset=self.gen_notebook_urn(notebook),
-                    type=DatasetLineageTypeClass.TRANSFORMED,
+                self._create_upstream_class(
+                    self.gen_notebook_urn(notebook.id),
+                    DatasetLineageTypeClass.TRANSFORMED,
+                    timestamp,
                 )
             )
 
@@ -771,6 +799,31 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             instance=self.config.platform_instance,
         ).as_urn()
 
+    def _create_upstream_class(
+        self,
+        dataset_urn: str,
+        lineage_type: Union[str, DatasetLineageTypeClass],
+        timestamp: Optional[int],
+    ) -> UpstreamClass:
+        """
+        Helper method to create UpstreamClass with optional audit stamp.
+        If timestamp is None, audit stamp is omitted.
+        """
+        if timestamp is not None:
+            return UpstreamClass(
+                dataset=dataset_urn,
+                type=lineage_type,
+                auditStamp=AuditStampClass(
+                    time=timestamp,
+                    actor=UNKNOWN_USER,
+                ),
+            )
+        else:
+            return UpstreamClass(
+                dataset=dataset_urn,
+                type=lineage_type,
+            )
+
     def gen_schema_containers(self, schema: Schema) -> Iterable[MetadataWorkUnit]:
         domain_urn = self._gen_domain_urn(f"{schema.catalog.name}.{schema.name}")
         schema_tags = []
@@ -961,16 +1014,20 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
         created: Optional[TimeStampClass] = None
         if table.created_at:
             custom_properties["created_at"] = str(table.created_at)
-            created = TimeStampClass(
-                int(table.created_at.timestamp() * 1000),
-                make_user_urn(table.created_by) if table.created_by else None,
-            )
+            created_ts = make_ts_millis(table.created_at)
+            if created_ts is not None:
+                created = TimeStampClass(
+                    created_ts,
+                    make_user_urn(table.created_by) if table.created_by else None,
+                )
         last_modified = created
         if table.updated_at:
-            last_modified = TimeStampClass(
-                int(table.updated_at.timestamp() * 1000),
-                table.updated_by and make_user_urn(table.updated_by),
-            )
+            updated_ts = make_ts_millis(table.updated_at)
+            if updated_ts is not None:
+                last_modified = TimeStampClass(
+                    updated_ts,
+                    table.updated_by and make_user_urn(table.updated_by),
+                )
 
         return DatasetPropertiesClass(
             name=table.name,

datahub/ingestion/source/unity/tag_entities.py

@@ -36,7 +36,7 @@ class UnityCatalogTagPlatformResourceId(BaseModel, ExternalEntityId):
 
     tag_key: str
     tag_value: Optional[str] = None
-    platform_instance: Optional[str]
+    platform_instance: Optional[str] = None
     exists_in_unity_catalog: bool = False
     persisted: bool = False
 
@@ -218,7 +218,7 @@ class UnityCatalogTagPlatformResource(BaseModel, ExternalEntity):
     datahub_urns: LinkedResourceSet
     managed_by_datahub: bool
     id: UnityCatalogTagPlatformResourceId
-    allowed_values: Optional[List[str]]
+    allowed_values: Optional[List[str]] = None
 
     def get_id(self) -> ExternalEntityId:
         return self.id

datahub/ingestion/source/usage/starburst_trino_usage.py

@@ -60,7 +60,7 @@ AggregatedDataset = GenericAggregatedDataset[TrinoTableRef]
 
 class TrinoConnectorInfo(BaseModel):
     partitionIds: List[str]
-    truncated: Optional[bool]
+    truncated: Optional[bool] = None
 
 
 class TrinoAccessedMetadata(BaseModel):
@@ -80,7 +80,7 @@ class TrinoJoinedAccessEvent(BaseModel):
     table: Optional[str] = None
     accessed_metadata: List[TrinoAccessedMetadata]
     starttime: datetime = Field(alias="create_time")
-    endtime: Optional[datetime] = Field(alias="end_time")
+    endtime: Optional[datetime] = Field(None, alias="end_time")
 
 
 class EnvBasedSourceBaseConfig:

datahub/ingestion/transformer/base_transformer.py

@@ -281,11 +281,14 @@ class BaseTransformer(Transformer, metaclass=ABCMeta):
                                 )
                             )
 
-                            record_metadata = _update_work_unit_id(
-                                envelope=envelope,
-                                aspect_name=mcp.aspect.get_aspect_name(),  # type: ignore
-                                urn=mcp.entityUrn,
-                            )
+                            if mcp.entityUrn:
+                                record_metadata = _update_work_unit_id(
+                                    envelope=envelope,
+                                    aspect_name=mcp.aspect.get_aspect_name(),  # type: ignore
+                                    urn=mcp.entityUrn,
+                                )
+                            else:
+                                record_metadata = envelope.metadata.copy()
 
                             yield RecordEnvelope(
                                 record=mcp,

datahub/sdk/search_client.py

@@ -112,6 +112,8 @@ class SearchClient:
         self,
         query: Optional[str] = None,
         filter: Optional[Filter] = None,
+        *,
+        skip_cache: bool = False,
     ) -> Iterable[Urn]:
         # TODO: Add better limit / pagination support.
         types, compiled_filters = compile_filters(filter)
@@ -120,5 +122,6 @@ class SearchClient:
             status=None,
             extra_or_filters=compiled_filters,
             entity_types=types,
+            skip_cache=skip_cache,
         ):
             yield Urn.from_string(urn)

datahub/specific/aspect_helpers/fine_grained_lineage.py

@@ -0,0 +1,76 @@
+from abc import abstractmethod
+from typing import List, Tuple
+
+from typing_extensions import Self
+
+from datahub.emitter.mcp_patch_builder import MetadataPatchProposal, PatchPath
+from datahub.metadata.schema_classes import (
+    FineGrainedLineageClass as FineGrainedLineage,
+)
+
+
+class HasFineGrainedLineagePatch(MetadataPatchProposal):
+    @abstractmethod
+    def _fine_grained_lineage_location(self) -> Tuple[str, PatchPath]:
+        """Return the aspect name where fine-grained lineage is stored."""
+        raise NotImplementedError("Subclasses must implement this method.")
+
+    @staticmethod
+    def _get_fine_grained_key(
+        fine_grained_lineage: FineGrainedLineage,
+    ) -> Tuple[str, str, str]:
+        downstreams = fine_grained_lineage.downstreams or []
+        if len(downstreams) != 1:
+            raise TypeError("Cannot patch with more or less than one downstream.")
+        transform_op = fine_grained_lineage.transformOperation or "NONE"
+        downstream_urn = downstreams[0]
+        query_id = fine_grained_lineage.query or "NONE"
+        return transform_op, downstream_urn, query_id
+
+    def add_fine_grained_lineage(
+        self, fine_grained_lineage: FineGrainedLineage
+    ) -> Self:
+        aspect_name, path = self._fine_grained_lineage_location()
+        (
+            transform_op,
+            downstream_urn,
+            query_id,
+        ) = self._get_fine_grained_key(fine_grained_lineage)
+        for upstream_urn in fine_grained_lineage.upstreams or []:
+            self._add_patch(
+                aspect_name,
+                "add",
+                path=(*path, transform_op, downstream_urn, query_id, upstream_urn),
+                value={"confidenceScore": fine_grained_lineage.confidenceScore},
+            )
+        return self
+
+    def remove_fine_grained_lineage(
+        self, fine_grained_lineage: FineGrainedLineage
+    ) -> Self:
+        aspect_name, path = self._fine_grained_lineage_location()
+        (
+            transform_op,
+            downstream_urn,
+            query_id,
+        ) = self._get_fine_grained_key(fine_grained_lineage)
+        for upstream_urn in fine_grained_lineage.upstreams or []:
+            self._add_patch(
+                aspect_name,
+                "remove",
+                path=(*path, transform_op, downstream_urn, query_id, upstream_urn),
+                value={},
+            )
+        return self
+
+    def set_fine_grained_lineages(
+        self, fine_grained_lineages: List[FineGrainedLineage]
+    ) -> Self:
+        aspect_name, path = self._fine_grained_lineage_location()
+        self._add_patch(
+            aspect_name,
+            "add",
+            path=path,
+            value=fine_grained_lineages,
+        )
+        return self

datahub/specific/datajob.py

@@ -1,15 +1,19 @@
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Set, Tuple, Union
 
 from datahub.emitter.mcp_patch_builder import MetadataPatchProposal, PatchPath
 from datahub.metadata.schema_classes import (
     DataJobInfoClass as DataJobInfo,
     DataJobInputOutputClass as DataJobInputOutput,
     EdgeClass as Edge,
+    FineGrainedLineageClass as FineGrainedLineage,
     KafkaAuditHeaderClass,
     SystemMetadataClass,
 )
 from datahub.metadata.urns import SchemaFieldUrn, Urn
 from datahub.specific.aspect_helpers.custom_properties import HasCustomPropertiesPatch
+from datahub.specific.aspect_helpers.fine_grained_lineage import (
+    HasFineGrainedLineagePatch,
+)
 from datahub.specific.aspect_helpers.ownership import HasOwnershipPatch
 from datahub.specific.aspect_helpers.tags import HasTagsPatch
 from datahub.specific.aspect_helpers.terms import HasTermsPatch
@@ -20,6 +24,7 @@ class DataJobPatchBuilder(
     HasCustomPropertiesPatch,
     HasTagsPatch,
     HasTermsPatch,
+    HasFineGrainedLineagePatch,
     MetadataPatchProposal,
 ):
     def __init__(
@@ -40,10 +45,19 @@ class DataJobPatchBuilder(
             urn, system_metadata=system_metadata, audit_header=audit_header
         )
 
+        # Track fine-grained lineages for DataJob-specific handling
+        self._fine_grained_lineages_to_add: List[FineGrainedLineage] = []
+        self._fine_grained_lineage_keys_to_remove: Set[Tuple[str, str, str]] = set()
+        self._fine_grained_lineages_set: Optional[List[FineGrainedLineage]] = None
+
     @classmethod
     def _custom_properties_location(cls) -> Tuple[str, PatchPath]:
         return DataJobInfo.ASPECT_NAME, ("customProperties",)
 
+    @classmethod
+    def _fine_grained_lineage_location(cls) -> Tuple[str, PatchPath]:
+        return DataJobInputOutput.ASPECT_NAME, ("fineGrainedLineages",)
+
     def add_input_datajob(self, input: Union[Edge, Urn, str]) -> "DataJobPatchBuilder":
         """
         Adds an input data job to the DataJobPatchBuilder.

datahub/specific/dataset.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Generic, List, Optional, Tuple, TypeVar, Union
 
 from datahub.emitter.mcp_patch_builder import MetadataPatchProposal, PatchPath
@@ -17,6 +18,9 @@ from datahub.metadata.schema_classes import (
 )
 from datahub.metadata.urns import DatasetUrn, TagUrn, Urn
 from datahub.specific.aspect_helpers.custom_properties import HasCustomPropertiesPatch
+from datahub.specific.aspect_helpers.fine_grained_lineage import (
+    HasFineGrainedLineagePatch,
+)
 from datahub.specific.aspect_helpers.ownership import HasOwnershipPatch
 from datahub.specific.aspect_helpers.structured_properties import (
     HasStructuredPropertiesPatch,
@@ -99,6 +103,7 @@ class DatasetPatchBuilder(
     HasStructuredPropertiesPatch,
     HasTagsPatch,
     HasTermsPatch,
+    HasFineGrainedLineagePatch,
     MetadataPatchProposal,
 ):
     def __init__(
@@ -115,6 +120,10 @@ class DatasetPatchBuilder(
     def _custom_properties_location(cls) -> Tuple[str, PatchPath]:
         return DatasetProperties.ASPECT_NAME, ("customProperties",)
 
+    @classmethod
+    def _fine_grained_lineage_location(cls) -> Tuple[str, PatchPath]:
+        return UpstreamLineage.ASPECT_NAME, ("fineGrainedLineages",)
+
     def add_upstream_lineage(self, upstream: Upstream) -> "DatasetPatchBuilder":
         self._add_patch(
             UpstreamLineage.ASPECT_NAME,
@@ -144,75 +153,44 @@ class DatasetPatchBuilder(
     def add_fine_grained_upstream_lineage(
         self, fine_grained_lineage: FineGrainedLineage
     ) -> "DatasetPatchBuilder":
-        (
-            transform_op,
-            downstream_urn,
-            query_id,
-        ) = DatasetPatchBuilder.get_fine_grained_key(fine_grained_lineage)
-        for upstream_urn in fine_grained_lineage.upstreams or []:
-            self._add_patch(
-                UpstreamLineage.ASPECT_NAME,
-                "add",
-                path=self._build_fine_grained_path(
-                    transform_op, downstream_urn, query_id, upstream_urn
-                ),
-                value={"confidenceScore": fine_grained_lineage.confidenceScore},
-            )
-        return self
-
-    @staticmethod
-    def get_fine_grained_key(
-        fine_grained_lineage: FineGrainedLineage,
-    ) -> Tuple[str, str, str]:
-        downstreams = fine_grained_lineage.downstreams or []
-        if len(downstreams) != 1:
-            raise TypeError("Cannot patch with more or less than one downstream.")
-        transform_op = fine_grained_lineage.transformOperation or "NONE"
-        downstream_urn = downstreams[0]
-        query_id = fine_grained_lineage.query or "NONE"
-        return transform_op, downstream_urn, query_id
-
-    @classmethod
-    def _build_fine_grained_path(
-        cls, transform_op: str, downstream_urn: str, query_id: str, upstream_urn: str
-    ) -> PatchPath:
-        return (
-            "fineGrainedLineages",
-            transform_op,
-            downstream_urn,
-            query_id,
-            upstream_urn,
+        """
+        Deprecated: Use `add_fine_grained_lineage` instead.
+        """
+        warnings.warn(
+            "add_fine_grained_upstream_lineage() is deprecated."
+            " Use add_fine_grained_lineage() instead.",
+            DeprecationWarning,
+            stacklevel=2,
         )
+        return self.add_fine_grained_lineage(fine_grained_lineage)
 
     def remove_fine_grained_upstream_lineage(
         self, fine_grained_lineage: FineGrainedLineage
     ) -> "DatasetPatchBuilder":
-        (
-            transform_op,
-            downstream_urn,
-            query_id,
-        ) = DatasetPatchBuilder.get_fine_grained_key(fine_grained_lineage)
-        for upstream_urn in fine_grained_lineage.upstreams or []:
-            self._add_patch(
-                UpstreamLineage.ASPECT_NAME,
-                "remove",
-                path=self._build_fine_grained_path(
-                    transform_op, downstream_urn, query_id, upstream_urn
-                ),
-                value={},
-            )
-        return self
+        """
+        Deprecated: Use `remove_fine_grained_lineage` instead.
+        """
+        warnings.warn(
+            "remove_fine_grained_upstream_lineage() is deprecated."
+            " Use remove_fine_grained_lineage() instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.remove_fine_grained_lineage(fine_grained_lineage)
 
     def set_fine_grained_upstream_lineages(
         self, fine_grained_lineages: List[FineGrainedLineage]
     ) -> "DatasetPatchBuilder":
-        self._add_patch(
-            UpstreamLineage.ASPECT_NAME,
-            "add",
-            path=("fineGrainedLineages",),
-            value=fine_grained_lineages,
+        """
+        Deprecated: Use `set_fine_grained_lineages` instead.
+        """
+        warnings.warn(
+            "set_fine_grained_upstream_lineages() is deprecated."
+            " Use set_fine_grained_lineages() instead.",
+            DeprecationWarning,
+            stacklevel=2,
         )
-        return self
+        return self.set_fine_grained_lineages(fine_grained_lineages)
 
     def for_field(
         self, field_path: str, editable: bool = True

datahub/utilities/server_config_util.py

@@ -242,7 +242,8 @@ class RestServiceConfig:
 
         # Check if this is a config-based feature
         if feature in config_based_features:
-            return config_based_features[feature]()
+            result = config_based_features[feature]()
+            return bool(result) if result is not None else False
 
         # For environment-based features, determine requirements based on cloud vs. non-cloud
         deployment_type = "cloud" if self.is_datahub_cloud else "core"