acryl-datahub 1.2.0.11rc4__py3-none-any.whl → 1.3.0__py3-none-any.whl

datahub/ingestion/api/source.py

@@ -25,7 +25,6 @@ from typing_extensions import LiteralString, Self
 
 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 @@ from datahub.ingestion.api.source_helpers import (
     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 @@ class StructuredLogs(Report):
 
 
 @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 @@ class Source(Closeable, metaclass=ABCMeta):
     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."
         )

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."""
+        ...

datahub/ingestion/source/bigquery_v2/bigquery_report.py

@@ -9,7 +9,6 @@ import pydantic
 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,
 ):

datahub/ingestion/source/bigquery_v2/bigquery_schema_gen.py

@@ -66,7 +66,7 @@ from datahub.ingestion.source.sql.sql_utils import (
 )
 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 @@ class BigQuerySchemaGenerator:
 
         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,

datahub/ingestion/source/cassandra/cassandra_profiling.py

@@ -18,7 +18,7 @@ from datahub.ingestion.source.cassandra.cassandra_api import (
 )
 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 @@ class CassandraProfiler:
         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,

datahub/ingestion/source/cassandra/cassandra_utils.py

@@ -6,7 +6,6 @@ from datahub.ingestion.source.cassandra.cassandra_api import CassandraColumn
 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 @@ SYSTEM_KEYSPACE_LIST = set(
 
 
 @dataclass
-class CassandraSourceReport(StaleEntityRemovalSourceReport, IngestionStageReport):
+class CassandraSourceReport(StaleEntityRemovalSourceReport):
     num_tables_failed: int = 0
     num_views_failed: int = 0
     tables_scanned: int = 0

datahub/ingestion/source/dremio/dremio_reporting.py

@@ -6,7 +6,6 @@ from datahub.ingestion.source.sql.sql_report import SQLSourceReport
 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 @@ from datahub.utilities.stats_collections import (
 class DremioSourceReport(
     SQLSourceReport,
     StaleEntityRemovalSourceReport,
-    IngestionStageReport,
     BaseTimeWindowReport,
 ):
     num_containers_failed: int = 0

datahub/ingestion/source/dremio/dremio_source.py

@@ -55,7 +55,7 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
 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 @@ class DremioSource(StatefulIngestionSourceBase):
             # 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,

datahub/ingestion/source/fivetran/config.py

@@ -29,6 +29,7 @@ from datahub.ingestion.source.state.stale_entity_removal_handler import (
 from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionConfigBase,
 )
+from datahub.ingestion.source.unity.connection import UnityCatalogConnectionConfig
 from datahub.utilities.lossy_collections import LossyList
 from datahub.utilities.perf_timer import PerfTimer
 
@@ -56,8 +57,8 @@ class Constant:
     STATUS = "status"
     USER_ID = "user_id"
     EMAIL = "email"
-    CONNECTOR_ID = "connector_id"
-    CONNECTOR_NAME = "connector_name"
+    CONNECTOR_ID = "connection_id"
+    CONNECTOR_NAME = "connection_name"
     CONNECTOR_TYPE_ID = "connector_type_id"
     PAUSED = "paused"
     SYNC_FREQUENCY = "sync_frequency"
@@ -85,10 +86,23 @@ class BigQueryDestinationConfig(BigQueryConnectionConfig):
     dataset: str = Field(description="The fivetran connector log dataset.")
 
 
+class DatabricksDestinationConfig(UnityCatalogConnectionConfig):
+    catalog: str = Field(description="The fivetran connector log catalog.")
+    log_schema: str = Field(description="The fivetran connector log schema.")
+
+    @pydantic.validator("warehouse_id")
+    def warehouse_id_should_not_be_empty(cls, warehouse_id: Optional[str]) -> str:
+        if warehouse_id is None or (warehouse_id and warehouse_id.strip() == ""):
+            raise ValueError("Fivetran requires warehouse_id to be set")
+        return warehouse_id
+
+
 class FivetranLogConfig(ConfigModel):
-    destination_platform: Literal["snowflake", "bigquery"] = pydantic.Field(
-        default="snowflake",
-        description="The destination platform where fivetran connector log tables are dumped.",
+    destination_platform: Literal["snowflake", "bigquery", "databricks"] = (
+        pydantic.Field(
+            default="snowflake",
+            description="The destination platform where fivetran connector log tables are dumped.",
+        )
     )
     snowflake_destination_config: Optional[SnowflakeDestinationConfig] = pydantic.Field(
         default=None,
@@ -98,6 +112,12 @@ class FivetranLogConfig(ConfigModel):
         default=None,
         description="If destination platform is 'bigquery', provide bigquery configuration.",
     )
+    databricks_destination_config: Optional[DatabricksDestinationConfig] = (
+        pydantic.Field(
+            default=None,
+            description="If destination platform is 'databricks', provide databricks configuration.",
+        )
+    )
     _rename_destination_config = pydantic_renamed_field(
         "destination_config", "snowflake_destination_config"
     )
@@ -115,6 +135,11 @@ class FivetranLogConfig(ConfigModel):
                 raise ValueError(
                     "If destination platform is 'bigquery', user must provide bigquery destination configuration in the recipe."
                 )
+        elif destination_platform == "databricks":
+            if "databricks_destination_config" not in values:
+                raise ValueError(
+                    "If destination platform is 'databricks', user must provide databricks destination configuration in the recipe."
+                )
         else:
             raise ValueError(
                 f"Destination platform '{destination_platform}' is not yet supported."

datahub/ingestion/source/fivetran/fivetran.py

@@ -66,7 +66,6 @@ logger = logging.getLogger(__name__)
 class FivetranSource(StatefulIngestionSourceBase):
     """
     This plugin extracts fivetran users, connectors, destinations and sync history.
-    This plugin is in beta and has only been tested on Snowflake connector.
     """
 
     config: FivetranSourceConfig

datahub/ingestion/source/fivetran/fivetran_log_api.py

@@ -73,6 +73,19 @@ class FivetranLogAPI:
                 if result is None:
                     raise ValueError("Failed to retrieve BigQuery project ID")
                 fivetran_log_database = result[0]
+        elif destination_platform == "databricks":
+            databricks_destination_config = (
+                self.fivetran_log_config.databricks_destination_config
+            )
+            if databricks_destination_config is not None:
+                engine = create_engine(
+                    databricks_destination_config.get_sql_alchemy_url(
+                        databricks_destination_config.catalog
+                    ),
+                    **databricks_destination_config.get_options(),
+                )
+                fivetran_log_query.set_schema(databricks_destination_config.log_schema)
+                fivetran_log_database = databricks_destination_config.catalog
         else:
             raise ConfigurationError(
                 f"Destination platform '{destination_platform}' is not yet supported."

datahub/ingestion/source/fivetran/fivetran_query.py

@@ -6,6 +6,21 @@ MAX_COLUMN_LINEAGE_PER_CONNECTOR = 1000
 MAX_JOBS_PER_CONNECTOR = 500
 
 
+"""
+------------------------------------------------------------------------------------------------------------
+Fivetran Platform Connector Handling
+------------------------------------------------------------------------------------------------------------
+Current Query Change Log: August 2025 (See: https://fivetran.com/docs/changelog/2025/august-2025)
+
+All queries have to be updated as per Fivetran Platform Connector release if any. We expect customers
+and fivetran to keep platform connector configured for DataHub with auto sync enabled to get latest changes.
+
+References:
+- Fivetran Release Notes: https://fivetran.com/docs/changelog (Look for "Fivetran Platform Connector")
+- Latest Platform Connector Schema: https://fivetran.com/docs/logs/fivetran-platform?erdModal=open
+"""
+
+
 class FivetranLogQuery:
     # Note: All queries are written in Snowflake SQL.
     # They will be transpiled to the target database's SQL dialect at runtime.
@@ -30,17 +45,17 @@ class FivetranLogQuery:
     def get_connectors_query(self) -> str:
         return f"""\
 SELECT
-  connector_id,
+  connection_id,
   connecting_user_id,
   connector_type_id,
-  connector_name,
+  connection_name,
   paused,
   sync_frequency,
   destination_id
-FROM {self.schema_clause}connector
+FROM {self.schema_clause}connection
 WHERE
   _fivetran_deleted = FALSE
-QUALIFY ROW_NUMBER() OVER (PARTITION BY connector_id ORDER BY _fivetran_synced DESC) = 1
+QUALIFY ROW_NUMBER() OVER (PARTITION BY connection_id ORDER BY _fivetran_synced DESC) = 1
 """
 
     def get_users_query(self) -> str:
@@ -63,20 +78,20 @@ FROM {self.schema_clause}user
         return f"""\
 WITH ranked_syncs AS (
     SELECT
-        connector_id,
+        connection_id,
         sync_id,
         MAX(CASE WHEN message_event = 'sync_start' THEN time_stamp END) as start_time,
         MAX(CASE WHEN message_event = 'sync_end' THEN time_stamp END) as end_time,
         MAX(CASE WHEN message_event = 'sync_end' THEN message_data END) as end_message_data,
-        ROW_NUMBER() OVER (PARTITION BY connector_id ORDER BY MAX(time_stamp) DESC) as rn
+        ROW_NUMBER() OVER (PARTITION BY connection_id ORDER BY MAX(time_stamp) DESC) as rn
     FROM {self.schema_clause}log
     WHERE message_event in ('sync_start', 'sync_end')
     AND time_stamp > CURRENT_TIMESTAMP - INTERVAL '{syncs_interval} days'
-    AND connector_id IN ({formatted_connector_ids})
-    GROUP BY connector_id, sync_id
+    AND connection_id IN ({formatted_connector_ids})
+    GROUP BY connection_id, sync_id
 )
 SELECT
-    connector_id,
+    connection_id,
     sync_id,
     start_time,
     end_time,
@@ -85,7 +100,7 @@ FROM ranked_syncs
 WHERE rn <= {MAX_JOBS_PER_CONNECTOR}
     AND start_time IS NOT NULL
     AND end_time IS NOT NULL
-ORDER BY connector_id, end_time DESC
+ORDER BY connection_id, end_time DESC
 """
 
     def get_table_lineage_query(self, connector_ids: List[str]) -> str:
@@ -97,7 +112,7 @@ SELECT
     *
 FROM (
     SELECT
-        stm.connector_id as connector_id,
+        stm.connection_id as connection_id,
         stm.id as source_table_id,
         stm.name as source_table_name,
         ssm.name as source_schema_name,
@@ -105,18 +120,18 @@ FROM (
         dtm.name as destination_table_name,
         dsm.name as destination_schema_name,
         tl.created_at as created_at,
-        ROW_NUMBER() OVER (PARTITION BY stm.connector_id, stm.id, dtm.id ORDER BY tl.created_at DESC) as table_combo_rn
+        ROW_NUMBER() OVER (PARTITION BY stm.connection_id, stm.id, dtm.id ORDER BY tl.created_at DESC) as table_combo_rn
     FROM {self.schema_clause}table_lineage as tl
-    JOIN {self.schema_clause}source_table_metadata as stm on tl.source_table_id = stm.id
-    JOIN {self.schema_clause}destination_table_metadata as dtm on tl.destination_table_id = dtm.id
-    JOIN {self.schema_clause}source_schema_metadata as ssm on stm.schema_id = ssm.id
-    JOIN {self.schema_clause}destination_schema_metadata as dsm on dtm.schema_id = dsm.id
-    WHERE stm.connector_id IN ({formatted_connector_ids})
+    JOIN {self.schema_clause}source_table as stm on tl.source_table_id = stm.id -- stm: source_table_metadata
+    JOIN {self.schema_clause}destination_table as dtm on tl.destination_table_id = dtm.id -- dtm: destination_table_metadata
+    JOIN {self.schema_clause}source_schema as ssm on stm.schema_id = ssm.id -- ssm: source_schema_metadata
+    JOIN {self.schema_clause}destination_schema as dsm on dtm.schema_id = dsm.id -- dsm: destination_schema_metadata
+    WHERE stm.connection_id IN ({formatted_connector_ids})
 )
 -- Ensure that we only get back one entry per source and destination pair.
 WHERE table_combo_rn = 1
-QUALIFY ROW_NUMBER() OVER (PARTITION BY connector_id ORDER BY created_at DESC) <= {MAX_TABLE_LINEAGE_PER_CONNECTOR}
-ORDER BY connector_id, created_at DESC
+QUALIFY ROW_NUMBER() OVER (PARTITION BY connection_id ORDER BY created_at DESC) <= {MAX_TABLE_LINEAGE_PER_CONNECTOR}
+ORDER BY connection_id, created_at DESC
 """
 
     def get_column_lineage_query(self, connector_ids: List[str]) -> str:
@@ -131,25 +146,25 @@ SELECT
     destination_column_name
 FROM (
     SELECT
-        stm.connector_id as connector_id,
+        stm.connection_id as connection_id,
         scm.table_id as source_table_id,
         dcm.table_id as destination_table_id,
         scm.name as source_column_name,
         dcm.name as destination_column_name,
         cl.created_at as created_at,
-        ROW_NUMBER() OVER (PARTITION BY stm.connector_id, cl.source_column_id, cl.destination_column_id ORDER BY cl.created_at DESC) as column_combo_rn
+        ROW_NUMBER() OVER (PARTITION BY stm.connection_id, cl.source_column_id, cl.destination_column_id ORDER BY cl.created_at DESC) as column_combo_rn
     FROM {self.schema_clause}column_lineage as cl
-    JOIN {self.schema_clause}source_column_metadata as scm
+    JOIN {self.schema_clause}source_column as scm -- scm: source_column_metadata
       ON cl.source_column_id = scm.id
-    JOIN {self.schema_clause}destination_column_metadata as dcm
+    JOIN {self.schema_clause}destination_column as dcm -- dcm: destination_column_metadata
       ON cl.destination_column_id = dcm.id
-    -- Only joining source_table_metadata to get the connector_id.
-    JOIN {self.schema_clause}source_table_metadata as stm
+    -- Only joining source_table to get the connection_id.
+    JOIN {self.schema_clause}source_table as stm -- stm: source_table_metadata
       ON scm.table_id = stm.id
-    WHERE stm.connector_id IN ({formatted_connector_ids})
+    WHERE stm.connection_id IN ({formatted_connector_ids})
 )
 -- Ensure that we only get back one entry per (connector, source column, destination column) pair.
 WHERE column_combo_rn = 1
-QUALIFY ROW_NUMBER() OVER (PARTITION BY connector_id ORDER BY created_at DESC) <= {MAX_COLUMN_LINEAGE_PER_CONNECTOR}
-ORDER BY connector_id, created_at DESC
+QUALIFY ROW_NUMBER() OVER (PARTITION BY connection_id ORDER BY created_at DESC) <= {MAX_COLUMN_LINEAGE_PER_CONNECTOR}
+ORDER BY connection_id, created_at DESC
 """

datahub/ingestion/source/gc/datahub_gc.py

@@ -34,7 +34,6 @@ from datahub.ingestion.source.gc.soft_deleted_entity_cleanup import (
     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
 

datahub/ingestion/source/grafana/models.py

@@ -8,12 +8,14 @@ References:
 - Dashboard JSON structure: https://grafana.com/docs/grafana/latest/dashboards/build-dashboards/view-dashboard-json-model/
 """
 
+import logging
 from typing import Any, Dict, List, Optional
 
 from pydantic import BaseModel, ConfigDict, Field
 
 from datahub.emitter.mcp_builder import ContainerKey
 
+logger = logging.getLogger(__name__)
 # Grafana-specific type definitions for better type safety
 GrafanaQueryTarget = Dict[
     str, Any
@@ -89,7 +91,13 @@ class Dashboard(_GrafanaBaseModel):
     def parse_obj(cls, data: Dict[str, Any]) -> "Dashboard":
         """Custom parsing to handle nested panel extraction."""
         dashboard_data = data.get("dashboard", {})
-        panels = cls.extract_panels(dashboard_data.get("panels", []))
+        _panel_data = dashboard_data.get("panels", [])
+        try:
+            panels = cls.extract_panels(_panel_data)
+        except Exception as e:
+            logger.warning(
+                f"Error extracting panels from dashboard for dashboard panels {_panel_data} : {e}"
+            )
 
         # Extract meta.folderId from nested structure
         meta = dashboard_data.get("meta", {})

datahub/ingestion/source/grafana/report.py

@@ -3,11 +3,10 @@ from dataclasses import dataclass
 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

datahub/ingestion/source/hex/hex.py

@@ -46,7 +46,6 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionConfigBase,
     StatefulIngestionSourceBase,
 )
-from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
 from datahub.sdk.main_client import DataHubClient
 
 
@@ -172,7 +171,6 @@ class HexSourceConfig(
 class HexReport(
     StaleEntityRemovalSourceReport,
     HexApiReport,
-    IngestionStageReport,
     HexQueryFetcherReport,
 ):
     pass

datahub/ingestion/source/redshift/redshift.py

@@ -89,8 +89,8 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
 from datahub.ingestion.source_report.ingestion_stage import (
     LINEAGE_EXTRACTION,
     METADATA_EXTRACTION,
-    PROFILING,
     USAGE_EXTRACTION_INGESTION,
+    IngestionHighStage,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.common import SubTypes, TimeStamp
 from datahub.metadata.com.linkedin.pegasus2avro.dataset import (
@@ -446,7 +446,7 @@ class RedshiftSource(StatefulIngestionSourceBase, TestableSource):
                 )
 
         if self.config.is_profiling_enabled():
-            with self.report.new_stage(PROFILING):
+            with self.report.new_high_stage(IngestionHighStage.PROFILING):
                 profiler = RedshiftProfiler(
                     config=self.config,
                     report=self.report,

datahub/ingestion/source/redshift/report.py

@@ -4,7 +4,6 @@ from typing import Dict, Optional
 
 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
@@ -15,7 +14,6 @@ from datahub.utilities.stats_collections import TopKDict
 @dataclass
 class RedshiftReport(
     SQLSourceReport,
-    IngestionStageReport,
     BaseTimeWindowReport,
     ClassificationReportMixin,
 ):

datahub/ingestion/source/snowflake/snowflake_report.py

@@ -9,7 +9,6 @@ from datahub.ingestion.source.sql.sql_report import SQLSourceReport
 from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionReport,
 )
-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
@@ -96,7 +95,6 @@ class SnowflakeV2Report(
     SnowflakeUsageReport,
     StatefulIngestionReport,
     ClassificationReportMixin,
-    IngestionStageReport,
 ):
     account_locator: Optional[str] = None
     region: Optional[str] = None

datahub/ingestion/source/snowflake/snowflake_schema_gen.py

@@ -77,7 +77,7 @@ from datahub.ingestion.source_report.ingestion_stage import (
     EXTERNAL_TABLE_DDL_LINEAGE,
     LINEAGE_EXTRACTION,
     METADATA_EXTRACTION,
-    PROFILING,
+    IngestionHighStage,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.common import (
     GlobalTags,
@@ -360,7 +360,7 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
         yield from self._process_db_schemas(snowflake_db, db_tables)
 
         if self.profiler and db_tables:
-            with self.report.new_stage(f"{snowflake_db.name}: {PROFILING}"):
+            with self.report.new_high_stage(IngestionHighStage.PROFILING):
                 yield from self.profiler.get_workunits(snowflake_db, db_tables)
 
     def _process_db_schemas(

datahub/ingestion/source/sql/oracle.py

@@ -37,7 +37,7 @@ from datahub.ingestion.source.sql.sql_config import (
 
 logger = logging.getLogger(__name__)
 
-oracledb.version = "8.3.0"
+oracledb.version = "8.3.0"  # type: ignore[assignment]
 sys.modules["cx_Oracle"] = oracledb
 
 extra_oracle_types = {