feat(ingest): add high level stage for ingestion (#14862)

Author: anshbansal

metadata-ingestion/src/datahub/ingestion/api/source.py

@@ -25,7 +25,6 @@
 
 from datahub.configuration.common import ConfigModel
 from datahub.configuration.source_common import PlatformInstanceConfigMixin
-from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.auto_work_units.auto_dataset_properties_aspect import (
     auto_patch_last_modified,
 )
@@ -46,8 +45,15 @@
     auto_workunit,
     auto_workunit_reporter,
 )
+from datahub.ingestion.api.source_protocols import (
+    MetadataWorkUnitIterable,
+    ProfilingCapable,
+)
 from datahub.ingestion.api.workunit import MetadataWorkUnit
-from datahub.sdk.entity import Entity
+from datahub.ingestion.source_report.ingestion_stage import (
+    IngestionHighStage,
+    IngestionStageReport,
+)
 from datahub.telemetry import stats
 from datahub.utilities.lossy_collections import LossyDict, LossyList
 from datahub.utilities.type_annotations import get_class_from_annotation
@@ -205,7 +211,7 @@ def infos(self) -> LossyList[StructuredLogEntry]:
 
 
 @dataclass
-class SourceReport(ExamplesReport):
+class SourceReport(ExamplesReport, IngestionStageReport):
     event_not_produced_warn: bool = True
     events_produced: int = 0
     events_produced_per_sec: int = 0
@@ -553,13 +559,31 @@ def _apply_workunit_processors(
     def get_workunits(self) -> Iterable[MetadataWorkUnit]:
         workunit_processors = self.get_workunit_processors()
         workunit_processors.append(AutoSystemMetadata(self.ctx).stamp)
-        return self._apply_workunit_processors(
+        # Process main workunits
+        yield from self._apply_workunit_processors(
             workunit_processors, auto_workunit(self.get_workunits_internal())
         )
+        # Process profiling workunits
+        yield from self._process_profiling_stage(workunit_processors)
+
+    def _process_profiling_stage(
+        self, processors: List[Optional[MetadataWorkUnitProcessor]]
+    ) -> Iterable[MetadataWorkUnit]:
+        """Process profiling stage if source supports it."""
+        if (
+            not isinstance(self, ProfilingCapable)
+            or not self.is_profiling_enabled_internal()
+        ):
+            return
+        with self.get_report().new_high_stage(IngestionHighStage.PROFILING):
+            profiling_stream = self._apply_workunit_processors(
+                processors, auto_workunit(self.get_profiling_internal())
+            )
+            yield from profiling_stream
 
     def get_workunits_internal(
         self,
-    ) -> Iterable[Union[MetadataWorkUnit, MetadataChangeProposalWrapper, Entity]]:
+    ) -> MetadataWorkUnitIterable:
         raise NotImplementedError(
             "get_workunits_internal must be implemented if get_workunits is not overriden."
         )

metadata-ingestion/src/datahub/ingestion/api/source_protocols.py

@@ -0,0 +1,23 @@
+from typing import Iterable, Protocol, Union, runtime_checkable
+
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.sdk.entity import Entity
+
+# Type alias for metadata work units - Python 3.9 compatible
+MetadataWorkUnitIterable = Iterable[
+    Union[MetadataWorkUnit, MetadataChangeProposalWrapper, Entity]
+]
+
+
+@runtime_checkable
+class ProfilingCapable(Protocol):
+    """Protocol for sources that support profiling functionality."""
+
+    def is_profiling_enabled_internal(self) -> bool:
+        """Check if profiling is enabled for this source."""
+        ...
+
+    def get_profiling_internal(self) -> MetadataWorkUnitIterable:
+        """Generate profiling work units."""
+        ...

metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_report.py

@@ -9,7 +9,6 @@
 from datahub.ingestion.api.report import Report
 from datahub.ingestion.glossary.classification_mixin import ClassificationReportMixin
 from datahub.ingestion.source.sql.sql_report import SQLSourceReport
-from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
 from datahub.ingestion.source_report.time_window import BaseTimeWindowReport
 from datahub.sql_parsing.sql_parsing_aggregator import SqlAggregatorReport
 from datahub.utilities.lossy_collections import LossyDict, LossyList, LossySet
@@ -78,7 +77,6 @@ class BigQueryQueriesExtractorReport(Report):
 @dataclass
 class BigQueryV2Report(
     SQLSourceReport,
-    IngestionStageReport,
     BaseTimeWindowReport,
     ClassificationReportMixin,
 ):

metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_schema_gen.py

@@ -66,7 +66,7 @@
 )
 from datahub.ingestion.source_report.ingestion_stage import (
     METADATA_EXTRACTION,
-    PROFILING,
+    IngestionHighStage,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.common import (
     Status,
@@ -416,7 +416,7 @@ def _process_project(
 
         if self.config.is_profiling_enabled():
             logger.info(f"Starting profiling project {project_id}")
-            with self.report.new_stage(f"{project_id}: {PROFILING}"):
+            with self.report.new_high_stage(IngestionHighStage.PROFILING):
                 yield from self.profiler.get_workunits(
                     project_id=project_id,
                     tables=db_tables,

metadata-ingestion/src/datahub/ingestion/source/cassandra/cassandra_profiling.py

@@ -18,7 +18,7 @@
 )
 from datahub.ingestion.source.cassandra.cassandra_config import CassandraSourceConfig
 from datahub.ingestion.source.cassandra.cassandra_utils import CassandraSourceReport
-from datahub.ingestion.source_report.ingestion_stage import PROFILING
+from datahub.ingestion.source_report.ingestion_stage import IngestionHighStage
 from datahub.metadata.schema_classes import (
     DatasetFieldProfileClass,
     DatasetProfileClass,
@@ -71,7 +71,7 @@ def get_workunits(
         for keyspace_name in cassandra_data.keyspaces:
             tables = cassandra_data.tables.get(keyspace_name, [])
             with (
-                self.report.new_stage(f"{keyspace_name}: {PROFILING}"),
+                self.report.new_high_stage(IngestionHighStage.PROFILING),
                 ThreadPoolExecutor(
                     max_workers=self.config.profiling.max_workers
                 ) as executor,

metadata-ingestion/src/datahub/ingestion/source/cassandra/cassandra_utils.py

@@ -6,7 +6,6 @@
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalSourceReport,
 )
-from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
 from datahub.metadata.com.linkedin.pegasus2avro.schema import (
     SchemaField,
     SchemaFieldDataType,
@@ -35,7 +34,7 @@
 
 
 @dataclass
-class CassandraSourceReport(StaleEntityRemovalSourceReport, IngestionStageReport):
+class CassandraSourceReport(StaleEntityRemovalSourceReport):
     num_tables_failed: int = 0
     num_views_failed: int = 0
     tables_scanned: int = 0

metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_reporting.py

@@ -6,7 +6,6 @@
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalSourceReport,
 )
-from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
 from datahub.ingestion.source_report.time_window import BaseTimeWindowReport
 from datahub.sql_parsing.sql_parsing_aggregator import SqlAggregatorReport
 from datahub.utilities.stats_collections import (
@@ -20,7 +19,6 @@
 class DremioSourceReport(
     SQLSourceReport,
     StaleEntityRemovalSourceReport,
-    IngestionStageReport,
     BaseTimeWindowReport,
 ):
     num_containers_failed: int = 0

metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_source.py

@@ -55,7 +55,7 @@
 from datahub.ingestion.source_report.ingestion_stage import (
     LINEAGE_EXTRACTION,
     METADATA_EXTRACTION,
-    PROFILING,
+    IngestionHighStage,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.dataset import (
     DatasetLineageTypeClass,
@@ -283,7 +283,7 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
             # Profiling
             if self.config.is_profiling_enabled():
                 with (
-                    self.report.new_stage(PROFILING),
+                    self.report.new_high_stage(IngestionHighStage.PROFILING),
                     ThreadPoolExecutor(
                         max_workers=self.config.profiling.max_workers
                     ) as executor,

metadata-ingestion/src/datahub/ingestion/source/gc/datahub_gc.py

@@ -34,7 +34,6 @@
     SoftDeletedEntitiesCleanupConfig,
     SoftDeletedEntitiesReport,
 )
-from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
 
 logger = logging.getLogger(__name__)
 
@@ -87,7 +86,6 @@ class DataHubGcSourceReport(
     DataProcessCleanupReport,
     SoftDeletedEntitiesReport,
     DatahubExecutionRequestCleanupReport,
-    IngestionStageReport,
 ):
     expired_tokens_revoked: int = 0
 

metadata-ingestion/src/datahub/ingestion/source/grafana/report.py

@@ -3,11 +3,10 @@
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalSourceReport,
 )
-from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
 
 
 @dataclass
-class GrafanaSourceReport(StaleEntityRemovalSourceReport, IngestionStageReport):
+class GrafanaSourceReport(StaleEntityRemovalSourceReport):
     # Entity counters
     dashboards_scanned: int = 0
     charts_scanned: int = 0