acryl-datahub-cloud 0.3.14.1rc5__tar.gz → 0.3.15rc0__tar.gz

{acryl_datahub_cloud-0.3.14.1rc5 → acryl_datahub_cloud-0.3.15rc0}/PKG-INFO

@@ -1,101 +1,98 @@
 Metadata-Version: 2.4
 Name: acryl-datahub-cloud
-Version: 0.3.14.1rc5
+Version: 0.3.15rc0
 Requires-Python: >=3.10
 Requires-Dist: avro-gen3==0.7.16
 Requires-Dist: acryl-datahub
+Requires-Dist: acryl-datahub==1.3.0
 Requires-Dist: croniter
 Requires-Dist: pytz
 Requires-Dist: types-croniter
 Requires-Dist: tzlocal
 Provides-Extra: datahub-lineage-features
 Requires-Dist: tenacity; extra == "datahub-lineage-features"
-Requires-Dist: pyarrow; extra == "datahub-lineage-features"
-Requires-Dist: opensearch-py==2.4.2; extra == "datahub-lineage-features"
 Requires-Dist: duckdb; extra == "datahub-lineage-features"
 Requires-Dist: pandas; extra == "datahub-lineage-features"
-Requires-Dist: pydantic<2; extra == "datahub-lineage-features"
+Requires-Dist: opensearch-py==2.4.2; extra == "datahub-lineage-features"
+Requires-Dist: pyarrow; extra == "datahub-lineage-features"
 Provides-Extra: datahub-reporting-forms
-Requires-Dist: pyarrow; extra == "datahub-reporting-forms"
-Requires-Dist: boto3; extra == "datahub-reporting-forms"
-Requires-Dist: termcolor==2.5.0; extra == "datahub-reporting-forms"
 Requires-Dist: duckdb; extra == "datahub-reporting-forms"
 Requires-Dist: pandas; extra == "datahub-reporting-forms"
-Requires-Dist: pydantic<2; extra == "datahub-reporting-forms"
+Requires-Dist: termcolor==2.5.0; extra == "datahub-reporting-forms"
+Requires-Dist: pyarrow; extra == "datahub-reporting-forms"
+Requires-Dist: boto3; extra == "datahub-reporting-forms"
 Provides-Extra: datahub-reporting-extract-graph
-Requires-Dist: pyarrow; extra == "datahub-reporting-extract-graph"
-Requires-Dist: boto3; extra == "datahub-reporting-extract-graph"
-Requires-Dist: opensearch-py==2.4.2; extra == "datahub-reporting-extract-graph"
 Requires-Dist: duckdb; extra == "datahub-reporting-extract-graph"
 Requires-Dist: pandas; extra == "datahub-reporting-extract-graph"
-Requires-Dist: pydantic<2; extra == "datahub-reporting-extract-graph"
+Requires-Dist: opensearch-py==2.4.2; extra == "datahub-reporting-extract-graph"
+Requires-Dist: pyarrow; extra == "datahub-reporting-extract-graph"
+Requires-Dist: boto3; extra == "datahub-reporting-extract-graph"
 Provides-Extra: datahub-reporting-extract-sql
+Requires-Dist: pandas; extra == "datahub-reporting-extract-sql"
+Requires-Dist: duckdb; extra == "datahub-reporting-extract-sql"
 Requires-Dist: pyarrow; extra == "datahub-reporting-extract-sql"
 Requires-Dist: boto3; extra == "datahub-reporting-extract-sql"
-Requires-Dist: duckdb; extra == "datahub-reporting-extract-sql"
-Requires-Dist: pandas; extra == "datahub-reporting-extract-sql"
-Requires-Dist: pydantic<2; extra == "datahub-reporting-extract-sql"
 Provides-Extra: datahub-usage-reporting
+Requires-Dist: polars==1.34.0; extra == "datahub-usage-reporting"
+Requires-Dist: scipy<=1.14.1; extra == "datahub-usage-reporting"
 Requires-Dist: pyarrow; extra == "datahub-usage-reporting"
-Requires-Dist: boto3; extra == "datahub-usage-reporting"
-Requires-Dist: opensearch-py==2.4.2; extra == "datahub-usage-reporting"
-Requires-Dist: duckdb; extra == "datahub-usage-reporting"
 Requires-Dist: numpy<2; extra == "datahub-usage-reporting"
-Requires-Dist: termcolor==2.5.0; extra == "datahub-usage-reporting"
-Requires-Dist: scipy<=1.14.1; extra == "datahub-usage-reporting"
-Requires-Dist: polars==1.30.0; extra == "datahub-usage-reporting"
 Requires-Dist: pyarrow<=18.0.0; extra == "datahub-usage-reporting"
-Requires-Dist: pandas; extra == "datahub-usage-reporting"
+Requires-Dist: duckdb; extra == "datahub-usage-reporting"
 Requires-Dist: elasticsearch==7.13.4; extra == "datahub-usage-reporting"
-Requires-Dist: pydantic<2; extra == "datahub-usage-reporting"
+Requires-Dist: pandas; extra == "datahub-usage-reporting"
+Requires-Dist: termcolor==2.5.0; extra == "datahub-usage-reporting"
+Requires-Dist: opensearch-py==2.4.2; extra == "datahub-usage-reporting"
+Requires-Dist: boto3; extra == "datahub-usage-reporting"
 Provides-Extra: datahub-metadata-sharing
 Requires-Dist: tenacity; extra == "datahub-metadata-sharing"
 Provides-Extra: datahub-action-request-owner
 Requires-Dist: tenacity; extra == "datahub-action-request-owner"
 Provides-Extra: acryl-cs-issues
-Requires-Dist: openai; extra == "acryl-cs-issues"
 Requires-Dist: zenpy; extra == "acryl-cs-issues"
-Requires-Dist: jinja2; extra == "acryl-cs-issues"
 Requires-Dist: slack-sdk; extra == "acryl-cs-issues"
+Requires-Dist: openai; extra == "acryl-cs-issues"
+Requires-Dist: jinja2; extra == "acryl-cs-issues"
 Provides-Extra: datahub-forms-notifications
 Requires-Dist: tenacity; extra == "datahub-forms-notifications"
 Provides-Extra: all
-Requires-Dist: tenacity; extra == "all"
-Requires-Dist: boto3; extra == "all"
 Requires-Dist: slack-sdk; extra == "all"
-Requires-Dist: duckdb; extra == "all"
-Requires-Dist: numpy<2; extra == "all"
+Requires-Dist: polars==1.34.0; extra == "all"
 Requires-Dist: scipy<=1.14.1; extra == "all"
+Requires-Dist: pyarrow; extra == "all"
+Requires-Dist: numpy<2; extra == "all"
 Requires-Dist: pyarrow<=18.0.0; extra == "all"
+Requires-Dist: tenacity; extra == "all"
+Requires-Dist: duckdb; extra == "all"
 Requires-Dist: elasticsearch==7.13.4; extra == "all"
+Requires-Dist: jinja2; extra == "all"
 Requires-Dist: pandas; extra == "all"
-Requires-Dist: zenpy; extra == "all"
-Requires-Dist: pyarrow; extra == "all"
+Requires-Dist: termcolor==2.5.0; extra == "all"
 Requires-Dist: opensearch-py==2.4.2; extra == "all"
+Requires-Dist: zenpy; extra == "all"
 Requires-Dist: openai; extra == "all"
-Requires-Dist: jinja2; extra == "all"
-Requires-Dist: termcolor==2.5.0; extra == "all"
-Requires-Dist: polars==1.30.0; extra == "all"
-Requires-Dist: pydantic<2; extra == "all"
+Requires-Dist: boto3; extra == "all"
 Provides-Extra: dev
+Requires-Dist: numpy<2; extra == "dev"
+Requires-Dist: duckdb; extra == "dev"
+Requires-Dist: pyarrow-stubs; extra == "dev"
+Requires-Dist: elasticsearch==7.13.4; extra == "dev"
+Requires-Dist: openai; extra == "dev"
+Requires-Dist: pyarrow<=18.0.0; extra == "dev"
+Requires-Dist: pandas; extra == "dev"
+Requires-Dist: acryl-datahub[dev]; extra == "dev"
+Requires-Dist: boto3; extra == "dev"
+Requires-Dist: pandas-stubs; extra == "dev"
+Requires-Dist: pyarrow; extra == "dev"
 Requires-Dist: tenacity; extra == "dev"
 Requires-Dist: zenpy; extra == "dev"
-Requires-Dist: pyarrow; extra == "dev"
-Requires-Dist: boto3; extra == "dev"
 Requires-Dist: opensearch-py==2.4.2; extra == "dev"
+Requires-Dist: scipy-stubs; extra == "dev"
 Requires-Dist: slack-sdk; extra == "dev"
-Requires-Dist: openai; extra == "dev"
-Requires-Dist: duckdb; extra == "dev"
-Requires-Dist: numpy<2; extra == "dev"
-Requires-Dist: elasticsearch==7.13.4; extra == "dev"
+Requires-Dist: scipy<=1.14.1; extra == "dev"
 Requires-Dist: jinja2; extra == "dev"
+Requires-Dist: polars==1.34.0; extra == "dev"
 Requires-Dist: termcolor==2.5.0; extra == "dev"
-Requires-Dist: scipy<=1.14.1; extra == "dev"
-Requires-Dist: polars==1.30.0; extra == "dev"
-Requires-Dist: pyarrow<=18.0.0; extra == "dev"
-Requires-Dist: pandas; extra == "dev"
-Requires-Dist: acryl-datahub[dev]; extra == "dev"
-Requires-Dist: pydantic<2; extra == "dev"
 Dynamic: provides-extra
 Dynamic: requires-dist
 Dynamic: requires-python

{acryl_datahub_cloud-0.3.14.1rc5 → acryl_datahub_cloud-0.3.15rc0}/setup.py

@@ -6,17 +6,22 @@ from setuptools import setup
 _codegen_config_file = pathlib.Path("./src/acryl_datahub_cloud/_codegen_config.json")
 _codegen_config: dict = json.loads(_codegen_config_file.read_text())
 
+acryl_datahub = [
+    # 1.3.0: Pydantic v2 was introduced in 1.2.0.10 and later RCs fixed in OSS some missing code from /metadata-ingestion SaaS
+    "acryl-datahub==1.3.0"
+]
+
 # Note: We are using the croniter library for cron parsing which is different from executor, which uses apscheduler, so there is a risk of mismatch here.
 # croniter is now maintained at: https://github.com/pallets-eco/croniter
 base_requirements = [
+    *acryl_datahub,
     "croniter",
     "pytz",
     "types-croniter",
     "tzlocal",
 ]
 
-# Adding pydantic<2 since we use pydantic models to map to pyarrow models and that is only compatible in pydantic v1
-stats_common = {"pandas", "pyarrow", "duckdb", "pydantic<2"}
+stats_common = {"pandas", "pyarrow", "duckdb"}
 aws_common = {"boto3"}
 open_search_common = {"opensearch-py==2.4.2"}
 
@@ -33,7 +38,7 @@ plugins = {
     | aws_common
     | {
         "opensearch-py==2.4.2",
-        "polars==1.30.0",
+        "polars==1.34.0",
         "elasticsearch==7.13.4",
         "numpy<2",
         "scipy<=1.14.1",
@@ -49,6 +54,10 @@ plugins = {
 dev_requirements = {
     # acryl-datahub[dev] pulls in more things than are strictly necessary, but it's fine.
     "acryl-datahub[dev]",
+    # Type stubs for external libraries
+    "pyarrow-stubs",
+    "scipy-stubs",
+    "pandas-stubs",
     *list(
         dependency
         for plugin in [

{acryl_datahub_cloud-0.3.14.1rc5 → acryl_datahub_cloud-0.3.15rc0}/src/acryl_datahub_cloud/_codegen_config.json

@@ -1,6 +1,6 @@
 {
   "name": "acryl-datahub-cloud",
-  "version": "0.3.14.1rc5",
+  "version": "0.3.15rc0",
   "install_requires": [
     "avro-gen3==0.7.16",
     "acryl-datahub"

{acryl_datahub_cloud-0.3.14.1rc5 → acryl_datahub_cloud-0.3.15rc0}/src/acryl_datahub_cloud/datahub_forms_notifications/forms_notifications_source.py

@@ -391,7 +391,7 @@ class DataHubFormsNotificationsSource(Source):
         user_urns = []
         group_urns = []
 
-        extra_fields = [f for f in DataHubDatasetSearchRow.__fields__]
+        extra_fields = [f for f in DataHubDatasetSearchRow.model_fields]
         results = self.graph.get_results_by_filter(
             extra_or_filters=self._get_incomplete_assets_for_form(form_urn, form.type),
             extra_source_fields=extra_fields,

{acryl_datahub_cloud-0.3.14.1rc5 → acryl_datahub_cloud-0.3.15rc0}/src/acryl_datahub_cloud/datahub_reporting/datahub_dataset.py

@@ -5,7 +5,7 @@ import pathlib
 import tempfile
 import time
 from enum import Enum
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple, Union, cast
 
 import boto3
 import duckdb
@@ -73,7 +73,9 @@ class FileStoreBackedDatasetConfig(ConfigModel):
     store_platform: str = "s3"
     file_name: str = "data"
     file_extension: str = "parquet"
-    file_compression: str = "snappy"
+    file_compression: Literal[
+        "gzip", "bz2", "brotli", "lz4", "zstd", "snappy", "none"
+    ] = "snappy"
     file_overwrite_existing: bool = True
     snapshot_partitioning_strategy: str = PartitioningStrategy.DATE
     generate_presigned_url: bool = True
@@ -119,9 +121,14 @@ class DataHubBasedS3Dataset:
         self.local_file_path: str = (
             config.file if config.file else self._initialize_local_file()
         )
-        self.file_writer = None
+        self.file_writer: Optional[pq.ParquetWriter] = None
         self.schema = (
-            pa.schema([(x.name, x.type) for x in self.dataset_metadata.schemaFields])
+            pa.schema(
+                [
+                    pa.field(x.name, BaseModelRow.string_to_pyarrow_type(x.type))
+                    for x in self.dataset_metadata.schemaFields
+                ]
+            )
             if self.dataset_metadata.schemaFields
             else None
         )
@@ -163,14 +170,28 @@ class DataHubBasedS3Dataset:
                     self.schema = row.arrow_schema()
                 else:
                     # hail mary: infer schema from the first row and cast everything to string
-                    self.schema = pa.schema([(key, pa.string()) for key in row])
+                    self.schema = pa.schema([pa.field(key, pa.string()) for key in row])
                     self.stringify_row = True
 
             self._initialize_local_file()
+            # Map compression names to PyArrow format (most are direct mappings)
+            compression_map = {
+                "gzip": "gzip",
+                "bz2": "brotli",  # PyArrow doesn't support bz2, use brotli
+                "brotli": "brotli",
+                "lz4": "lz4",
+                "zstd": "zstd",
+                "snappy": "snappy",
+                "none": "none",
+            }
+            compression = cast(
+                Literal["gzip", "bz2", "brotli", "lz4", "zstd", "snappy", "none"],
+                compression_map.get(self.config.file_compression, "snappy"),
+            )
             self.file_writer = pq.ParquetWriter(
                 self.local_file_path,
                 self.schema,
-                compression=self.config.file_compression,
+                compression=compression,
             )
         if isinstance(row, (BaseModel, BaseModelRow)):
             # for anything extending BaseModel, we want to use the dict representation
@@ -396,7 +417,9 @@ class DataHubBasedS3Dataset:
                 assert dataset_profiles.fieldProfiles is not None
                 dataset_profiles.fieldProfiles.append(field_profile)
             logger.info("Generated dataset profile")
-            schema_metadata = self._generate_schema_metadata(columns)
+            schema_metadata = self._generate_schema_metadata(
+                [(col[0], col[1]) for col in columns]
+            )
         return dataset_profiles, schema_metadata
 
     def register_dataset(

{acryl_datahub_cloud-0.3.14.1rc5 → acryl_datahub_cloud-0.3.15rc0}/src/acryl_datahub_cloud/datahub_reporting/datahub_form_reporting.py

@@ -306,7 +306,7 @@ class DataHubFormReportingData(FormData):
         on_asset_scanned: Optional[Callable[[str], Any]] = None,
         on_form_scanned: Optional[Callable[[str], Any]] = None,
     ) -> Iterable[FormReportingRow]:
-        extra_fields = [f for f in self.DataHubDatasetSearchRow.__fields__]
+        extra_fields = [f for f in self.DataHubDatasetSearchRow.model_fields]
         # TODO: Replace with the new search/filter SDK.
         result = self.graph.get_results_by_filter(
             extra_or_filters=self.get_form_existence_or_filters(),

{acryl_datahub_cloud-0.3.14.1rc5 → acryl_datahub_cloud-0.3.15rc0}/src/acryl_datahub_cloud/datahub_usage_reporting/usage_feature_reporter.py

@@ -42,7 +42,7 @@ from datahub.ingestion.api.decorators import (
     platform_name,
     support_status,
 )
-from datahub.ingestion.api.source import MetadataWorkUnitProcessor, SourceReport
+from datahub.ingestion.api.source import MetadataWorkUnitProcessor
 from datahub.ingestion.api.source_helpers import auto_workunit_reporter
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.graph.client import DatahubClientConfig
@@ -239,7 +239,7 @@ def exp_cdf(series: polars.Series) -> polars.Series:
 
 
 @dataclass
-class DatahubUsageFeatureReport(IngestionStageReport, StatefulIngestionReport):
+class DatahubUsageFeatureReport(StatefulIngestionReport, IngestionStageReport):
     dataset_platforms_count: Dict[str, int] = field(
         default_factory=lambda: defaultdict(lambda: 0)
     )
@@ -738,17 +738,20 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
                 return pa.dictionary(index_type=pa.int32(), value_type=pa.string())
             elif isinstance(polars_dtype, polars.Struct):
                 return pa.struct(
-                    {
-                        field.name: convert_dtype(field.dtype)
+                    [
+                        pa.field(field.name, convert_dtype(field.dtype))
                         for field in polars_dtype.fields
-                    }
+                    ]
                 )
             elif isinstance(polars_dtype, polars.List):
                 return pa.list_(convert_dtype(polars_dtype.inner))
             else:
                 raise ValueError(f"Unsupported Polars dtype: {polars_dtype}")
 
-        fields = [(name, convert_dtype(dtype)) for name, dtype in polars_schema.items()]
+        fields = [
+            pa.field(name, convert_dtype(dtype))
+            for name, dtype in polars_schema.items()
+        ]
         return pa.schema(fields)
 
     def batch_write_parquet(
@@ -971,26 +974,27 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
 
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
         if self.config.user_usage_enabled:
-            self.report.new_stage("generate user usage")
-            yield from self.generate_user_usage_mcps()
+            with self.report.new_stage("generate user usage"):
+                yield from self.generate_user_usage_mcps()
 
         if self.config.dataset_usage_enabled:
-            self.report.new_stage("generate dataset usage")
-            yield from self.generate_dataset_usage_mcps()
+            with self.report.new_stage("generate dataset usage"):
+                yield from self.generate_dataset_usage_mcps()
 
         if self.config.dashboard_usage_enabled:
-            self.report.new_stage("generate dashboard usage")
-            yield from self.generate_dashboard_usage_mcps()
+            with self.report.new_stage("generate dashboard usage"):
+                yield from self.generate_dashboard_usage_mcps()
 
         if self.config.chart_usage_enabled:
-            self.report.new_stage("generate chart usage")
-            yield from self.generate_chart_usage_mcps()
+            with self.report.new_stage("generate chart usage"):
+                yield from self.generate_chart_usage_mcps()
 
         if self.config.query_usage_enabled:
-            self.report.new_stage("generate query usage")
-            yield from self.generate_query_usage_mcps()
+            with self.report.new_stage("generate query usage"):
+                yield from self.generate_query_usage_mcps()
 
-        self.report.new_stage("end so time is calculated for last stage")
+        with self.report.new_stage("end so time is calculated for last stage"):
+            pass
 
     def generate_mcp_from_lazyframe(
         self, lazy_frame: polars.LazyFrame
@@ -2091,5 +2095,5 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
                 )
                 time.sleep(delay)
 
-    def get_report(self) -> SourceReport:
+    def get_report(self) -> "DatahubUsageFeatureReport":
         return self.report

{acryl_datahub_cloud-0.3.14.1rc5 → acryl_datahub_cloud-0.3.15rc0}/src/acryl_datahub_cloud/elasticsearch/graph_service.py

@@ -32,15 +32,31 @@ class BaseModelRow(BaseModel):
         else:
             raise ValueError(f"No mapping for type {type_}")
 
+    @staticmethod
+    def string_to_pyarrow_type(type_string: str) -> pa.DataType:
+        """Convert string representation back to pyarrow type by converting to Python type first."""
+        # Mapping of pyarrow string representations to Python types
+        type_mapping = {
+            "string": str,
+            "int64": int,
+            "float64": float,
+            "bool": bool,
+            "timestamp[ns]": datetime.datetime,
+            "date32[day]": datetime.date,
+        }
+
+        python_type = type_mapping.get(
+            type_string, str
+        )  # Default to str for unknown types
+        return BaseModelRow.pydantic_type_to_pyarrow(python_type)
+
     @classmethod
     def arrow_schema(cls) -> pa.Schema:
         fields = []
-        for field_name, field_model in cls.__fields__.items():
-            pyarrow_type = BaseModelRow.pydantic_type_to_pyarrow(
-                field_model.outer_type_
-            )
+        for field_name, field_model in cls.model_fields.items():
+            pyarrow_type = BaseModelRow.pydantic_type_to_pyarrow(field_model.annotation)
             pyarrow_field = pa.field(field_name, pyarrow_type)
-            if not field_model.required:
+            if not field_model.is_required():
                 pyarrow_field = pyarrow_field.with_nullable(True)
             else:
                 pyarrow_field = pyarrow_field.with_nullable(False)
@@ -50,10 +66,8 @@ class BaseModelRow(BaseModel):
     @classmethod
     def datahub_schema(cls) -> List[SchemaField]:
         fields = []
-        for field_name, field_model in cls.__fields__.items():
-            pyarrow_type = BaseModelRow.pydantic_type_to_pyarrow(
-                field_model.outer_type_
-            )
+        for field_name, field_model in cls.model_fields.items():
+            pyarrow_type = BaseModelRow.pydantic_type_to_pyarrow(field_model.annotation)
             fields.append(SchemaField(name=field_name, type=str(pyarrow_type)))
         return fields
 

{acryl_datahub_cloud-0.3.14.1rc5 → acryl_datahub_cloud-0.3.15rc0}/src/acryl_datahub_cloud/lineage_features/source.py

@@ -3,7 +3,6 @@ import os
 import time
 from collections import defaultdict
 from dataclasses import dataclass
-from datetime import datetime, timezone
 from typing import Any, Callable, Dict, Iterable, List, Set
 
 from opensearchpy import OpenSearch
@@ -53,6 +52,12 @@ class LineageFeaturesSourceConfig(ConfigModel):
     retry_delay_seconds: int = 5
     retry_backoff_multiplier: float = 2.0
 
+    # Cleanup old features when they have not been updated for this many days
+    # This is required because we only emit this feature for cases where we find a lineage
+    # in the graph index
+    cleanup_batch_size: int = 100
+    cleanup_old_features_days: int = 2
+
     @validator("max_retries")
     def validate_max_retries(cls, v: int) -> int:
         if v < 1:
@@ -79,6 +84,12 @@ class LineageExtractGraphSourceReport(SourceReport, IngestionStageReport):
     downstream_count: int = 0
     edges_scanned: int = 0
     skipped_materialized_urns_count: int = 0
+    zero_upstream_count: int = 0
+    zero_downstream_count: int = 0
+    has_asset_level_lineage_count: int = 0
+    zero_asset_level_lineage_count: int = 0
+    cleanup_old_features_time: int = 0
+    cleanup_old_features_count: int = 0
 
 
 @platform_name(id="datahub", platform_name="DataHub")
@@ -255,7 +266,6 @@ class DataHubLineageFeaturesSource(Source):
         with self.report.new_stage("Load valid URNs"):
             self.populate_valid_urns()
 
-        timestamp = datetime.now(tz=timezone.utc)
         server = self._create_opensearch_client_with_retry()
 
         query = {
@@ -326,7 +336,58 @@ class DataHubLineageFeaturesSource(Source):
         self._update_report()
         self._delete_pit_with_retry(server, pit)
 
-        self.report.new_stage("start emission of lineage features")
+        with self.report.new_stage("emission of lineage features"):
+            yield from self._emit_lineage_features()
+
+        with self.report.new_stage("cleanup old lineage features"):
+            yield from self._cleanup_old_features()
+
+    def _cleanup_old_features(self) -> Iterable[MetadataWorkUnit]:
+        """
+        This is required because we only emit this feature for cases where we find a lineage
+        in the graph index
+        """
+        cutoff_time = int(
+            (time.time() - (self.config.cleanup_old_features_days * 24 * 60 * 60))
+            * 1000
+        )
+        self.report.cleanup_old_features_time = cutoff_time
+
+        for urn in self.ctx.require_graph("Cleanup old features").get_urns_by_filter(
+            extraFilters=[
+                {
+                    "field": "hasAssetLevelLineageFeature",
+                    "negated": False,
+                    "condition": "EQUAL",
+                    "values": ["true"],
+                },
+                {
+                    "field": "lineageFeaturesComputedAt",
+                    "negated": False,
+                    "condition": "LESS_THAN",
+                    "values": [str(cutoff_time)],
+                },
+            ],
+            batch_size=self.config.cleanup_batch_size,
+        ):
+            # Emit lineage features with zero upstreams and downstreams for cleanup
+            wu = MetadataChangeProposalWrapper(
+                entityUrn=urn,
+                aspect=LineageFeaturesClass(
+                    upstreamCount=0,
+                    downstreamCount=0,
+                    hasAssetLevelLineage=False,
+                    computedAt=AuditStampClass(
+                        time=int(time.time() * 1000),
+                        actor=SYSTEM_ACTOR,
+                    ),
+                ),
+            ).as_workunit()
+            self.report.cleanup_old_features_count += 1
+            self.report.report_workunit(wu)
+            yield wu
+
+    def _emit_lineage_features(self) -> Iterable[MetadataWorkUnit]:
         # In Python 3.9, can be replaced by `self.self.upstream_counts.keys() | self.downstream_counts.keys()`
         for urn in set(self.upstream_counts.keys()).union(
             self.downstream_counts.keys()
@@ -337,21 +398,31 @@ class DataHubLineageFeaturesSource(Source):
             logger.debug(
                 f"{urn}: {self.upstream_counts[urn]}, {self.downstream_counts[urn]}"
             )
+            if self.upstream_counts[urn] == 0:
+                self.report.zero_upstream_count += 1
+            if self.downstream_counts[urn] == 0:
+                self.report.zero_downstream_count += 1
+            has_asset_level_lineage = (
+                self.upstream_counts[urn] > 0 or self.downstream_counts[urn] > 0
+            )
+            if has_asset_level_lineage:
+                self.report.has_asset_level_lineage_count += 1
+            else:
+                self.report.zero_asset_level_lineage_count += 1
             wu = MetadataChangeProposalWrapper(
                 entityUrn=urn,
                 aspect=LineageFeaturesClass(
                     upstreamCount=self.upstream_counts[urn],
                     downstreamCount=self.downstream_counts[urn],
+                    hasAssetLevelLineage=has_asset_level_lineage,
                     computedAt=AuditStampClass(
-                        time=int(timestamp.timestamp() * 1000),
+                        time=int(time.time() * 1000),
                         actor=SYSTEM_ACTOR,
                     ),
                 ),
             ).as_workunit()
             self.report.report_workunit(wu)
             yield wu
-        # So previous stage's calculations are done
-        self.report.new_stage("end emission of lineage features")
 
     def get_report(self) -> SourceReport:
         return self.report

{acryl_datahub_cloud-0.3.14.1rc5 → acryl_datahub_cloud-0.3.15rc0}/src/acryl_datahub_cloud/metadata/_urns/urn_defs.py

@@ -2391,6 +2391,62 @@ class ActionRequestUrn(_SpecificUrn):
     def id(self) -> str:
         return self._entity_ids[0]
 
+if TYPE_CHECKING:
+    from datahub.metadata.schema_classes import DataHubFileKeyClass
+
+class DataHubFileUrn(_SpecificUrn):
+    ENTITY_TYPE: ClassVar[Literal["dataHubFile"]] = "dataHubFile"
+    _URN_PARTS: ClassVar[int] = 1
+
+    def __init__(self, id: Union["DataHubFileUrn", str], *, _allow_coercion: bool = True) -> None:
+        if _allow_coercion:
+            # Field coercion logic (if any is required).
+            if isinstance(id, str):
+                if id.startswith('urn:li:'):
+                    try:
+                        id = DataHubFileUrn.from_string(id)
+                    except InvalidUrnError:
+                        raise InvalidUrnError(f'Expecting a DataHubFileUrn but got {id}')
+                else:
+                    id = UrnEncoder.encode_string(id)
+
+        # Validation logic.
+        if not id:
+            raise InvalidUrnError("DataHubFileUrn id cannot be empty")
+        if isinstance(id, DataHubFileUrn):
+            id = id.id
+        elif isinstance(id, Urn):
+            raise InvalidUrnError(f'Expecting a DataHubFileUrn but got {id}')
+        if UrnEncoder.contains_reserved_char(id):
+            raise InvalidUrnError(f'DataHubFileUrn id contains reserved characters')
+
+        super().__init__(self.ENTITY_TYPE, [id])
+
+    @classmethod
+    def _parse_ids(cls, entity_ids: List[str]) -> "DataHubFileUrn":
+        if len(entity_ids) != cls._URN_PARTS:
+            raise InvalidUrnError(f"DataHubFileUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
+        return cls(id=entity_ids[0], _allow_coercion=False)
+
+    @classmethod
+    def underlying_key_aspect_type(cls) -> Type["DataHubFileKeyClass"]:
+        from datahub.metadata.schema_classes import DataHubFileKeyClass
+
+        return DataHubFileKeyClass
+
+    def to_key_aspect(self) -> "DataHubFileKeyClass":
+        from datahub.metadata.schema_classes import DataHubFileKeyClass
+
+        return DataHubFileKeyClass(id=self.id)
+
+    @classmethod
+    def from_key_aspect(cls, key_aspect: "DataHubFileKeyClass") -> "DataHubFileUrn":
+        return cls(id=key_aspect.id)
+
+    @property
+    def id(self) -> str:
+        return self._entity_ids[0]
+
 if TYPE_CHECKING:
     from datahub.metadata.schema_classes import DataProcessInstanceKeyClass
 
@@ -3537,6 +3593,10 @@ class DataJobUrn(_SpecificUrn):
     def get_data_flow_urn(self) -> "DataFlowUrn":
         return DataFlowUrn.from_string(self.flow)
 
+    @property
+    def orchestrator(self) -> str:
+        return self.get_data_flow_urn().orchestrator
+
     @deprecated(reason="Use .job_id instead")
     def get_job_id(self) -> str:
         return self.job_id

{acryl_datahub_cloud-0.3.14.1rc5 → acryl_datahub_cloud-0.3.15rc0}/src/acryl_datahub_cloud/metadata/com/linkedin/pegasus2avro/execution/__init__.py

@@ -7,6 +7,7 @@
 # pylint: skip-file
 # fmt: off
 # isort: skip_file
+from .....schema_classes import ExecutionRequestArtifactsLocationClass
 from .....schema_classes import ExecutionRequestInputClass
 from .....schema_classes import ExecutionRequestResultClass
 from .....schema_classes import ExecutionRequestSignalClass
@@ -14,6 +15,7 @@ from .....schema_classes import ExecutionRequestSourceClass
 from .....schema_classes import StructuredExecutionReportClass
 
 
+ExecutionRequestArtifactsLocation = ExecutionRequestArtifactsLocationClass
 ExecutionRequestInput = ExecutionRequestInputClass
 ExecutionRequestResult = ExecutionRequestResultClass
 ExecutionRequestSignal = ExecutionRequestSignalClass