acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/api/report.py

@@ -2,18 +2,34 @@ import dataclasses
 import json
 import logging
 import pprint
-from dataclasses import dataclass
+from collections import defaultdict
+from dataclasses import dataclass, field
 from datetime import datetime, timedelta
 from enum import Enum
-from typing import Any, Optional, runtime_checkable
+from typing import Any, Dict, List, Optional, Set, Union, cast, runtime_checkable
 
 import humanfriendly
 import pydantic
 from pydantic import BaseModel
+from tabulate import tabulate
 from typing_extensions import Literal, Protocol
 
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.emitter.mcp_builder import mcps_from_mce
+from datahub.ingestion.api.closeable import Closeable
 from datahub.ingestion.api.report_helpers import format_datetime_relative
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.autogenerated.lineage_helper import is_lineage_aspect
+from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
+from datahub.metadata.schema_classes import (
+    MetadataChangeProposalClass,
+    StatusClass,
+    SubTypesClass,
+    UpstreamLineageClass,
+)
+from datahub.utilities.file_backed_collections import FileBackedDict
 from datahub.utilities.lossy_collections import LossyList
+from datahub.utilities.urns.urn import guess_platform_name
 
 logger = logging.getLogger(__name__)
 LogLevel = Literal["ERROR", "WARNING", "INFO", "DEBUG"]
@@ -26,6 +42,15 @@ class SupportsAsObj(Protocol):
 
 @dataclass
 class Report(SupportsAsObj):
+    def __post_init__(self) -> None:
+        self.platform: Optional[str] = None
+
+    def set_platform(self, platform: str) -> None:
+        self.platform = platform
+
+    def get_platform(self) -> Optional[str]:
+        return self.platform
+
     @staticmethod
     def to_str(some_val: Any) -> str:
         if isinstance(some_val, Enum):
@@ -82,7 +107,58 @@ class Report(SupportsAsObj):
         }
 
     def as_string(self) -> str:
-        return pprint.pformat(self.as_obj(), width=150, sort_dicts=False)
+        self_obj = self.as_obj()
+        _aspects_by_subtypes = self_obj.pop("aspects_by_subtypes", None)
+
+        # Format the main report data
+        result = pprint.pformat(self_obj, width=150, sort_dicts=False)
+
+        # Add aspects_by_subtypes table if it exists
+        if _aspects_by_subtypes:
+            result += "\n\nAspects by Subtypes:\n"
+            result += self._format_aspects_by_subtypes_table(_aspects_by_subtypes)
+
+        return result
+
+    def _format_aspects_by_subtypes_table(
+        self, aspects_by_subtypes: Dict[str, Dict[str, Dict[str, int]]]
+    ) -> str:
+        """Format aspects_by_subtypes data as a table with aspects as rows and entity/subtype as columns."""
+        if not aspects_by_subtypes:
+            return "No aspects by subtypes data available."
+
+        all_aspects: set[str] = {
+            aspect
+            for subtypes in aspects_by_subtypes.values()
+            for aspects in subtypes.values()
+            for aspect in aspects
+        }
+
+        aspect_rows = sorted(all_aspects)
+
+        entity_subtype_columns = []
+        for entity_type, subtypes in aspects_by_subtypes.items():
+            for subtype in subtypes:
+                entity_subtype_columns.append(f"{entity_type} ({subtype})")
+
+        entity_subtype_columns.sort()
+
+        headers = ["Aspect"] + entity_subtype_columns
+
+        table_data = [
+            [aspect]
+            + [
+                aspects.get(aspect, 0)
+                for subtypes in aspects_by_subtypes.values()
+                for aspects in subtypes.values()
+            ]
+            for aspect in aspect_rows
+        ]
+
+        if table_data:
+            return tabulate(table_data, headers=headers, tablefmt="grid")
+        else:
+            return "No aspects by subtypes data available."
 
     def as_json(self) -> str:
         return json.dumps(self.as_obj())
@@ -90,6 +166,15 @@ class Report(SupportsAsObj):
     # TODO add helper method for warning / failure status + counts?
 
 
+@dataclass
+class SourceReportSubtypes:
+    urn: str
+    entity_type: str
+    subType: str = field(default="unknown")
+    aspects: Dict[str, int] = field(default_factory=dict)
+    soft_deleted: bool = field(default=False)
+
+
 class ReportAttribute(BaseModel):
     severity: LogLevel = "DEBUG"
     help: Optional[str] = None
@@ -108,6 +193,299 @@ class ReportAttribute(BaseModel):
         logger.log(level=self.logger_sev, msg=msg, stacklevel=3)
 
 
+@dataclass
+class ExamplesReport(Report, Closeable):
+    aspects: Dict[str, Dict[str, int]] = field(
+        default_factory=lambda: defaultdict(lambda: defaultdict(int))
+    )
+    # This counts existence of aspects for each entity/subtype
+    # This is used for the UI to calculate %age of entities with the aspect
+    aspects_by_subtypes: Dict[str, Dict[str, Dict[str, int]]] = field(
+        default_factory=lambda: defaultdict(
+            lambda: defaultdict(lambda: defaultdict(int))
+        )
+    )
+    # This counts all aspects for each entity/subtype
+    aspects_by_subtypes_full_count: Dict[str, Dict[str, Dict[str, int]]] = field(
+        default_factory=lambda: defaultdict(
+            lambda: defaultdict(lambda: defaultdict(int))
+        )
+    )
+    samples: Dict[str, Dict[str, List[str]]] = field(
+        default_factory=lambda: defaultdict(lambda: defaultdict(list))
+    )
+    compute_stats_time_seconds: float = 0.0
+    _file_based_dict: Optional[FileBackedDict[SourceReportSubtypes]] = None
+
+    # We are adding this to make querying easier for fine-grained lineage
+    _fine_grained_lineage_special_case_name = "fineGrainedLineages"
+    _samples_to_add: int = 20
+    _lineage_aspects_seen: Set[str] = field(default_factory=set)
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self._file_based_dict = FileBackedDict(
+            tablename="urn_aspects",
+            extra_columns={
+                "urn": lambda val: val.urn,
+                "entityType": lambda val: val.entity_type,
+                "subTypes": lambda val: val.subType,
+                "aspects": lambda val: json.dumps(val.aspects),
+                "soft_deleted": lambda val: val.soft_deleted,
+            },
+        )
+
+    def close(self) -> None:
+        self.compute_stats()
+        if self._file_based_dict is not None:
+            self._file_based_dict.close()
+            self._file_based_dict = None
+
+    def _build_aspects_where_clause(self, aspects: List[str]) -> str:
+        """Build WHERE clause for matching any of the given aspects."""
+        if not aspects:
+            return ""
+
+        conditions = []
+        for aspect in aspects:
+            conditions.append(f"aspects LIKE '%{aspect}%'")
+
+        return " OR ".join(conditions)
+
+    def _collect_samples_by_subtype(self, where_clause: str, sample_key: str) -> None:
+        """Helper method to collect samples organized by subtype for a given where clause."""
+
+        subtype_query = f"""
+        SELECT DISTINCT subTypes
+        FROM urn_aspects 
+        WHERE {where_clause}
+        """
+        assert self._file_based_dict is not None
+        subtypes = set()
+        for row in self._file_based_dict.sql_query(subtype_query):
+            sub_type = row["subTypes"] or "unknown"
+            subtypes.add(sub_type)
+
+        for sub_type in subtypes:
+            query = f"""
+            SELECT urn
+            FROM urn_aspects 
+            WHERE {where_clause} AND subTypes = ?
+            limit {self._samples_to_add}
+            """
+
+            for row in self._file_based_dict.sql_query(query, (sub_type,)):
+                self.samples[sample_key][sub_type].append(row["urn"])
+
+    def _collect_samples_by_aspects(self, aspects: List[str], sample_key: str) -> None:
+        """Helper method to collect samples for entities that have any of the given aspects."""
+        if not aspects:
+            return
+
+        where_clause = self._build_aspects_where_clause(aspects)
+        self._collect_samples_by_subtype(where_clause, sample_key)
+
+    def _collect_samples_by_lineage_aspects(
+        self, aspects: List[str], sample_key: str
+    ) -> None:
+        """Helper method to collect samples for entities that have any of the given lineage aspects.
+
+        Lineage aspects are stored in JSON format and require quote escaping in LIKE clauses.
+        """
+        if not aspects:
+            return
+
+        lineage_conditions = []
+        for aspect in aspects:
+            lineage_conditions.append(f"aspects LIKE '%\"{aspect}\"%'")
+
+        where_clause = " OR ".join(lineage_conditions)
+        self._collect_samples_by_subtype(where_clause, sample_key)
+
+    def _collect_samples_with_all_conditions(self, sample_key: str) -> None:
+        """
+        Collect samples for entities that have lineage, profiling, and usage aspects.
+        These specific 3 cases are added here as these URNs will be shown in the UI. Subject to change in future.
+        """
+        if not self._lineage_aspects_seen:
+            return
+        assert self._file_based_dict is not None
+
+        # Build lineage conditions using the same logic as _collect_samples_by_lineage_aspects
+        lineage_conditions = []
+        for aspect in self._lineage_aspects_seen:
+            lineage_conditions.append(f"aspects LIKE '%\"{aspect}\"%'")
+        lineage_where_clause = " OR ".join(lineage_conditions)
+
+        # Build profiling conditions using the same logic as _collect_samples_by_aspects
+        profiling_where_clause = self._build_aspects_where_clause(["datasetProfile"])
+
+        # Build usage conditions using the same logic as _collect_samples_by_aspects
+        usage_where_clause = self._build_aspects_where_clause(
+            [
+                "datasetUsageStatistics",
+                "chartUsageStatistics",
+                "dashboardUsageStatistics",
+            ]
+        )
+
+        query = f"""
+        SELECT urn, subTypes
+        FROM urn_aspects
+        WHERE ({lineage_where_clause})
+        AND ({profiling_where_clause})
+        AND ({usage_where_clause})
+        limit {self._samples_to_add}
+        """
+
+        for row in self._file_based_dict.sql_query(query):
+            sub_type = row["subTypes"] or "unknown"
+            self.samples[sample_key][sub_type].append(row["urn"])
+
+    def _has_fine_grained_lineage(
+        self, mcp: Union[MetadataChangeProposalClass, MetadataChangeProposalWrapper]
+    ) -> bool:
+        if isinstance(mcp.aspect, UpstreamLineageClass):
+            upstream_lineage = cast(UpstreamLineageClass, mcp.aspect)
+            if upstream_lineage.fineGrainedLineages:
+                return True
+        return False
+
+    def _update_file_based_dict(
+        self,
+        urn: str,
+        entityType: str,
+        aspectName: str,
+        mcp: Union[MetadataChangeProposalClass, MetadataChangeProposalWrapper],
+    ) -> None:
+        platform_name = guess_platform_name(urn)
+        if platform_name != self.get_platform():
+            return
+        if is_lineage_aspect(entityType, aspectName):
+            self._lineage_aspects_seen.add(aspectName)
+        has_fine_grained_lineage = self._has_fine_grained_lineage(mcp)
+
+        sub_type = "unknown"
+        if isinstance(mcp.aspect, SubTypesClass):
+            sub_type = mcp.aspect.typeNames[0]
+
+        assert self._file_based_dict is not None
+        if urn in self._file_based_dict:
+            if sub_type != "unknown":
+                self._file_based_dict[urn].subType = sub_type
+            aspects_dict = self._file_based_dict[urn].aspects
+            if aspectName in aspects_dict:
+                aspects_dict[aspectName] += 1
+            else:
+                aspects_dict[aspectName] = 1
+            if has_fine_grained_lineage:
+                if self._fine_grained_lineage_special_case_name in aspects_dict:
+                    aspects_dict[self._fine_grained_lineage_special_case_name] += 1
+                else:
+                    aspects_dict[self._fine_grained_lineage_special_case_name] = 1
+            self._file_based_dict.mark_dirty(urn)
+        else:
+            aspects_dict = {aspectName: 1}
+            if has_fine_grained_lineage:
+                aspects_dict[self._fine_grained_lineage_special_case_name] = 1
+            self._file_based_dict[urn] = SourceReportSubtypes(
+                urn=urn,
+                entity_type=entityType,
+                subType=sub_type,
+                aspects=aspects_dict,
+            )
+        if (
+            isinstance(mcp.aspect, StatusClass)
+            and mcp is not None
+            and mcp.aspect is not None
+        ):
+            self._file_based_dict[urn].soft_deleted = mcp.aspect.removed
+            self._file_based_dict.mark_dirty(urn)
+
+    def _store_workunit_data(self, wu: MetadataWorkUnit) -> None:
+        urn = wu.get_urn()
+
+        if not isinstance(wu.metadata, MetadataChangeEvent):
+            mcps = [wu.metadata]
+        else:
+            mcps = list(mcps_from_mce(wu.metadata))
+
+        for mcp in mcps:
+            entityType = mcp.entityType
+            aspectName = mcp.aspectName
+
+            if aspectName is None:
+                continue
+
+            self._update_file_based_dict(urn, entityType, aspectName, mcp)
+
+    def compute_stats(self) -> None:
+        start_time = datetime.now()
+        if self._file_based_dict is None:
+            return
+
+        query = """
+        SELECT entityType, subTypes, aspects, count(*) as count
+        FROM urn_aspects 
+        WHERE soft_deleted = 0
+        GROUP BY entityType, subTypes, aspects
+        """
+
+        entity_subtype_aspect_counts: Dict[str, Dict[str, Dict[str, int]]] = (
+            defaultdict(lambda: defaultdict(lambda: defaultdict(int)))
+        )
+        entity_subtype_aspect_counts_exist: Dict[str, Dict[str, Dict[str, int]]] = (
+            defaultdict(lambda: defaultdict(lambda: defaultdict(int)))
+        )
+
+        for row in self._file_based_dict.sql_query(query):
+            entity_type = row["entityType"]
+            sub_type = row["subTypes"]
+            count = row["count"]
+            aspects_raw = row["aspects"] or "[]"
+
+            aspects = json.loads(aspects_raw)
+            for aspect, aspect_count in aspects.items():
+                entity_subtype_aspect_counts[entity_type][sub_type][aspect] += (
+                    aspect_count * count
+                )
+                entity_subtype_aspect_counts_exist[entity_type][sub_type][aspect] += (
+                    count
+                )
+
+        self.aspects.clear()
+        self.aspects_by_subtypes.clear()
+        self.aspects_by_subtypes_full_count.clear()
+        for entity_type, subtype_counts in entity_subtype_aspect_counts.items():
+            for sub_type, aspect_counts in subtype_counts.items():
+                for aspect, count in aspect_counts.items():
+                    self.aspects[entity_type][aspect] += count
+                self.aspects_by_subtypes_full_count[entity_type][sub_type] = dict(
+                    aspect_counts
+                )
+
+        for entity_type, subtype_counts in entity_subtype_aspect_counts_exist.items():
+            for sub_type, aspect_counts in subtype_counts.items():
+                self.aspects_by_subtypes[entity_type][sub_type] = dict(aspect_counts)
+
+        self.samples.clear()
+        self._collect_samples_by_aspects(["datasetProfile"], "profiling")
+        self._collect_samples_by_aspects(
+            [
+                "datasetUsageStatistics",
+                "chartUsageStatistics",
+                "dashboardUsageStatistics",
+            ],
+            "usage",
+        )
+        self._collect_samples_by_lineage_aspects(
+            list(self._lineage_aspects_seen), "lineage"
+        )
+        self._collect_samples_with_all_conditions("all_3")
+        end_time = datetime.now()
+        self.compute_stats_time_seconds += (end_time - start_time).total_seconds()
+
+
 class EntityFilterReport(ReportAttribute):
     type: str
 

datahub/ingestion/api/sink.py

@@ -1,7 +1,8 @@
 import datetime
+import logging
 from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass, field
-from typing import Any, Generic, Optional, Type, TypeVar, cast
+from typing import Any, Callable, Generic, List, Optional, Type, TypeVar, cast
 
 from typing_extensions import Self
 
@@ -12,6 +13,8 @@ from datahub.ingestion.api.report import Report
 from datahub.utilities.lossy_collections import LossyList
 from datahub.utilities.type_annotations import get_class_from_annotation
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass
 class SinkReport(Report):
@@ -89,6 +92,7 @@ class Sink(Generic[SinkConfig, SinkReportType], Closeable, metaclass=ABCMeta):
     ctx: PipelineContext
     config: SinkConfig
     report: SinkReportType
+    _pre_shutdown_callbacks: List[Callable[[], None]]
 
     @classmethod
     def get_config_class(cls) -> Type[SinkConfig]:
@@ -106,6 +110,7 @@ class Sink(Generic[SinkConfig, SinkReportType], Closeable, metaclass=ABCMeta):
         self.ctx = ctx
         self.config = config
         self.report = self.get_report_class()()
+        self._pre_shutdown_callbacks = []
 
         self.__post_init__()
 
@@ -144,8 +149,28 @@ class Sink(Generic[SinkConfig, SinkReportType], Closeable, metaclass=ABCMeta):
     def get_report(self) -> SinkReportType:
         return self.report
 
+    def register_pre_shutdown_callback(self, callback: Callable[[], None]) -> None:
+        """Register a callback to be executed before the sink shuts down.
+
+        This is useful for components that need to send final reports or cleanup
+        operations before the sink's resources are released.
+        """
+        self._pre_shutdown_callbacks.append(callback)
+
     def close(self) -> None:
-        pass
+        """Close the sink and clean up resources.
+
+        This method executes any registered pre-shutdown callbacks before
+        performing the actual shutdown. Subclasses should override this method
+        to provide sink-specific cleanup logic while calling super().close()
+        to ensure callbacks are executed.
+        """
+        # Execute pre-shutdown callbacks before shutdown
+        for callback in self._pre_shutdown_callbacks:
+            try:
+                callback()
+            except Exception as e:
+                logger.warning(f"Pre-shutdown callback failed: {e}", exc_info=True)
 
     def configured(self) -> str:
         """Override this method to output a human-readable and scrubbed version of the configured sink"""