acryl-datahub 1.1.0.5rc6__py3-none-any.whl → 1.1.0.5rc8__py3-none-any.whl

datahub/ingestion/api/report.py

@@ -6,19 +6,25 @@ from collections import defaultdict
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta
 from enum import Enum
-from typing import Any, Dict, Optional, Set, cast, runtime_checkable
+from typing import Any, Dict, List, Optional, Set, Union, cast, runtime_checkable
 
 import humanfriendly
 import pydantic
 from pydantic import BaseModel
 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 SubTypesClass, UpstreamLineageClass
+from datahub.metadata.schema_classes import (
+    MetadataChangeProposalClass,
+    SubTypesClass,
+    UpstreamLineageClass,
+)
 from datahub.utilities.file_backed_collections import FileBackedDict
 from datahub.utilities.lossy_collections import LossyList
 
@@ -125,8 +131,6 @@ class ReportAttribute(BaseModel):
 
 @dataclass
 class ExamplesReport(Report, Closeable):
-    _urns_seen: Set[str] = field(default_factory=set)
-    entities: Dict[str, list] = field(default_factory=lambda: defaultdict(LossyList))
     aspects: Dict[str, Dict[str, int]] = field(
         default_factory=lambda: defaultdict(lambda: defaultdict(int))
     )
@@ -135,11 +139,16 @@ class ExamplesReport(Report, Closeable):
             lambda: defaultdict(lambda: defaultdict(int))
         )
     )
-    aspect_urn_samples: Dict[str, Dict[str, LossyList[str]]] = field(
-        default_factory=lambda: defaultdict(lambda: defaultdict(LossyList))
+    samples: Dict[str, Dict[str, List[str]]] = field(
+        default_factory=lambda: defaultdict(lambda: defaultdict(list))
     )
     _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:
         self._file_based_dict = FileBackedDict(
             tablename="urn_aspects",
@@ -157,6 +166,151 @@ class ExamplesReport(Report, Closeable):
             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:
+        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
+            self._file_based_dict[urn].aspects.add(aspectName)
+            if has_fine_grained_lineage:
+                self._file_based_dict[urn].aspects.add(
+                    self._fine_grained_lineage_special_case_name
+                )
+            self._file_based_dict.mark_dirty(urn)
+        else:
+            self._file_based_dict[urn] = SourceReportSubtypes(
+                urn=urn,
+                entity_type=entityType,
+                subType=sub_type,
+                aspects={aspectName}
+                if not has_fine_grained_lineage
+                else {aspectName, self._fine_grained_lineage_special_case_name},
+            )
+
     def _store_workunit_data(self, wu: MetadataWorkUnit) -> None:
         urn = wu.get_urn()
 
@@ -169,41 +323,15 @@ class ExamplesReport(Report, Closeable):
             entityType = mcp.entityType
             aspectName = mcp.aspectName
 
-            if urn not in self._urns_seen:
-                self._urns_seen.add(urn)
-                self.entities[entityType].append(urn)
-
             if aspectName is None:
                 continue
-            self.aspects[entityType][aspectName] += 1
-            self.aspect_urn_samples[entityType][aspectName].append(urn)
-            sub_type = "unknown"
-            if isinstance(mcp.aspect, UpstreamLineageClass):
-                upstream_lineage = cast(UpstreamLineageClass, mcp.aspect)
-                if upstream_lineage.fineGrainedLineages:
-                    self.aspect_urn_samples[entityType]["fineGrainedLineages"].append(
-                        urn
-                    )
-                    self.aspects[entityType]["fineGrainedLineages"] += 1
-            elif 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
-                self._file_based_dict[urn].aspects.add(aspectName)
-                self._file_based_dict.mark_dirty(urn)
-            else:
-                self._file_based_dict[urn] = SourceReportSubtypes(
-                    urn=urn,
-                    entity_type=entityType,
-                    subType=sub_type,
-                    aspects={aspectName},
-                )
+
+            self._update_file_based_dict(urn, entityType, aspectName, mcp)
 
     def compute_stats(self) -> None:
         if self._file_based_dict is None:
             return
+
         query = """
         SELECT entityType, subTypes, aspects, count(*) as count
         FROM urn_aspects 
@@ -223,11 +351,31 @@ class ExamplesReport(Report, Closeable):
             for aspect in aspects:
                 entity_subtype_aspect_counts[entity_type][sub_type][aspect] += count
 
+        self.aspects.clear()
         self.aspects_by_subtypes.clear()
+        _aspects_seen: Set[str] = set()
         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
+                    _aspects_seen.add(aspect)
                 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")
+
 
 class EntityFilterReport(ReportAttribute):
     type: str