acryl-datahub-cloud 0.3.14.1rc4__py3-none-any.whl → 0.3.15rc0__py3-none-any.whl

acryl_datahub_cloud/_codegen_config.json

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

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/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/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/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/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/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/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/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

acryl_datahub_cloud/metadata/com/linkedin/pegasus2avro/file/__init__.py

@@ -0,0 +1,19 @@
+# mypy: ignore-errors
+# flake8: noqa
+
+# This file is autogenerated by /metadata-ingestion/scripts/avro_codegen.py
+# Do not modify manually!
+
+# pylint: skip-file
+# fmt: off
+# isort: skip_file
+from .....schema_classes import BucketStorageLocationClass
+from .....schema_classes import DataHubFileInfoClass
+from .....schema_classes import FileUploadScenarioClass
+
+
+BucketStorageLocation = BucketStorageLocationClass
+DataHubFileInfo = DataHubFileInfoClass
+FileUploadScenario = FileUploadScenarioClass
+
+# fmt: on

acryl_datahub_cloud/metadata/com/linkedin/pegasus2avro/metadata/key/__init__.py

@@ -22,6 +22,7 @@ from ......schema_classes import DataFlowKeyClass
 from ......schema_classes import DataHubAccessTokenKeyClass
 from ......schema_classes import DataHubActionKeyClass
 from ......schema_classes import DataHubConnectionKeyClass
+from ......schema_classes import DataHubFileKeyClass
 from ......schema_classes import DataHubIngestionSourceKeyClass
 from ......schema_classes import DataHubMetricCubeKeyClass
 from ......schema_classes import DataHubOpenAPISchemaKeyClass
@@ -92,6 +93,7 @@ DataFlowKey = DataFlowKeyClass
 DataHubAccessTokenKey = DataHubAccessTokenKeyClass
 DataHubActionKey = DataHubActionKeyClass
 DataHubConnectionKey = DataHubConnectionKeyClass
+DataHubFileKey = DataHubFileKeyClass
 DataHubIngestionSourceKey = DataHubIngestionSourceKeyClass
 DataHubMetricCubeKey = DataHubMetricCubeKeyClass
 DataHubOpenAPISchemaKey = DataHubOpenAPISchemaKeyClass

acryl_datahub_cloud/metadata/com/linkedin/pegasus2avro/role/__init__.py

@@ -8,11 +8,13 @@
 # fmt: off
 # isort: skip_file
 from .....schema_classes import ActorsClass
+from .....schema_classes import RoleGroupClass
 from .....schema_classes import RolePropertiesClass
 from .....schema_classes import RoleUserClass
 
 
 Actors = ActorsClass
+RoleGroup = RoleGroupClass
 RoleProperties = RolePropertiesClass
 RoleUser = RoleUserClass
 

acryl_datahub_cloud/metadata/com/linkedin/pegasus2avro/settings/global/__init__.py

@@ -23,6 +23,8 @@ from ......schema_classes import GlobalSettingsInfoClass
 from ......schema_classes import GlobalViewsSettingsClass
 from ......schema_classes import GlobalVisualSettingsClass
 from ......schema_classes import HelpLinkClass
+from ......schema_classes import OAuthProviderClass
+from ......schema_classes import OAuthSettingsClass
 from ......schema_classes import OidcSettingsClass
 from ......schema_classes import SlackIntegrationSettingsClass
 from ......schema_classes import SsoSettingsClass
@@ -47,6 +49,8 @@ GlobalSettingsInfo = GlobalSettingsInfoClass
 GlobalViewsSettings = GlobalViewsSettingsClass
 GlobalVisualSettings = GlobalVisualSettingsClass
 HelpLink = HelpLinkClass
+OAuthProvider = OAuthProviderClass
+OAuthSettings = OAuthSettingsClass
 OidcSettings = OidcSettingsClass
 SlackIntegrationSettings = SlackIntegrationSettingsClass
 SsoSettings = SsoSettingsClass