refactor(ingest): simplify stateful ingestion provider interface (#8104)

Co-authored-by: Tamas Nemeth <[email protected]>

Author: hsheth2

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

@@ -4,7 +4,7 @@
 
 from datahub.emitter.mce_builder import set_dataset_urn_to_lower
 from datahub.ingestion.api.committable import Committable
-from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph
+from datahub.ingestion.graph.client import DataHubGraph
 
 if TYPE_CHECKING:
     from datahub.ingestion.run.pipeline import PipelineConfig
@@ -43,22 +43,19 @@ class PipelineContext:
     def __init__(
         self,
         run_id: str,
-        datahub_api: Optional["DatahubClientConfig"] = None,
+        graph: Optional[DataHubGraph] = None,
         pipeline_name: Optional[str] = None,
         dry_run: bool = False,
         preview_mode: bool = False,
         pipeline_config: Optional["PipelineConfig"] = None,
     ) -> None:
         self.pipeline_config = pipeline_config
+        self.graph = graph
         self.run_id = run_id
         self.pipeline_name = pipeline_name
         self.dry_run_mode = dry_run
         self.preview_mode = preview_mode
         self.checkpointers: Dict[str, Committable] = {}
-        try:
-            self.graph = DataHubGraph(datahub_api) if datahub_api is not None else None
-        except Exception as e:
-            raise Exception(f"Failed to connect to DataHub: {e}") from e
 
         self._set_dataset_urn_to_lower_if_needed()
 

metadata-ingestion/src/datahub/ingestion/run/pipeline.py

@@ -27,6 +27,7 @@
 from datahub.ingestion.api.source import Extractor, Source
 from datahub.ingestion.api.transform import Transformer
 from datahub.ingestion.extractor.extractor_registry import extractor_registry
+from datahub.ingestion.graph.client import DataHubGraph
 from datahub.ingestion.reporting.reporting_provider_registry import (
     reporting_provider_registry,
 )
@@ -183,10 +184,15 @@ def __init__(
         self.last_time_printed = int(time.time())
         self.cli_report = CliReport()
 
+        self.graph = None
+        with _add_init_error_context("connect to DataHub"):
+            if self.config.datahub_api:
+                self.graph = DataHubGraph(self.config.datahub_api)
+
         with _add_init_error_context("set up framework context"):
             self.ctx = PipelineContext(
                 run_id=self.config.run_id,
-                datahub_api=self.config.datahub_api,
+                graph=self.graph,
                 pipeline_name=self.config.pipeline_name,
                 dry_run=dry_run,
                 preview_mode=preview_mode,

metadata-ingestion/src/datahub/ingestion/sink/datahub_rest.py

@@ -7,7 +7,7 @@
 from datetime import timedelta
 from enum import auto
 from threading import BoundedSemaphore
-from typing import Union, cast
+from typing import Union
 
 from datahub.cli.cli_utils import set_env_variables_override_config
 from datahub.configuration.common import (
@@ -126,8 +126,7 @@ def __post_init__(self) -> None:
 
     def handle_work_unit_start(self, workunit: WorkUnit) -> None:
         if isinstance(workunit, MetadataWorkUnit):
-            mwu: MetadataWorkUnit = cast(MetadataWorkUnit, workunit)
-            self.treat_errors_as_warnings = mwu.treat_errors_as_warnings
+            self.treat_errors_as_warnings = workunit.treat_errors_as_warnings
 
     def handle_work_unit_end(self, workunit: WorkUnit) -> None:
         pass

metadata-ingestion/src/datahub/ingestion/source/looker/looker_source.py

@@ -1331,6 +1331,3 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
 
     def get_report(self) -> SourceReport:
         return self.reporter
-
-    def close(self):
-        self.prepare_for_commit()

metadata-ingestion/src/datahub/ingestion/source/looker/lookml_source.py

@@ -2163,6 +2163,3 @@ def get_internal_workunits(self) -> Iterable[MetadataWorkUnit]:  # noqa: C901
 
     def get_report(self):
         return self.reporter
-
-    def close(self):
-        self.prepare_for_commit()

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

@@ -1228,7 +1228,7 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
                     # Because job_id is used as dictionary key, we have to set a new job_id
                     # Refer to https://github.com/datahub-project/datahub/blob/master/metadata-ingestion/src/datahub/ingestion/source/state/stateful_ingestion_base.py#L390
                     self.stale_entity_removal_handler.set_job_id(workspace.id)
-                    self.register_stateful_ingestion_usecase_handler(
+                    self.state_provider.register_stateful_ingestion_usecase_handler(
                         self.stale_entity_removal_handler
                     )
 

metadata-ingestion/src/datahub/ingestion/source/sql/vertica.py

@@ -1283,6 +1283,3 @@ def _get_owner_information(self, table: str, label: str) -> Optional[str]:
                 return each["owner_name"]
 
         return None
-
-    def close(self):
-        self.prepare_for_commit()

metadata-ingestion/src/datahub/ingestion/source/state/entity_removal_state.py

@@ -1,11 +1,9 @@
-from typing import Any, Dict, Iterable, List, Type
+from typing import Any, Dict, Iterable, List, Tuple, Type
 
 import pydantic
 
 from datahub.emitter.mce_builder import make_assertion_urn, make_container_urn
-from datahub.ingestion.source.state.stale_entity_removal_handler import (
-    StaleEntityCheckpointStateBase,
-)
+from datahub.ingestion.source.state.checkpoint import CheckpointStateBase
 from datahub.utilities.checkpoint_state_util import CheckpointStateUtil
 from datahub.utilities.dedup_list import deduplicate_list
 from datahub.utilities.urns.urn import guess_entity_type
@@ -56,7 +54,7 @@ def _validate_field_rename(cls: Type, values: dict) -> dict:
     return pydantic.root_validator(pre=True, allow_reuse=True)(_validate_field_rename)
 
 
-class GenericCheckpointState(StaleEntityCheckpointStateBase["GenericCheckpointState"]):
+class GenericCheckpointState(CheckpointStateBase):
     urns: List[str] = pydantic.Field(default_factory=list)
 
     # We store a bit of extra internal-only state so that we can keep the urns list deduplicated.
@@ -85,21 +83,34 @@ def __init__(self, **data: Any):  # type: ignore
         self.urns = deduplicate_list(self.urns)
         self._urns_set = set(self.urns)
 
-    @classmethod
-    def get_supported_types(cls) -> List[str]:
-        return ["*"]
-
     def add_checkpoint_urn(self, type: str, urn: str) -> None:
+        """
+        Adds an urn into the list used for tracking the type.
+
+        :param type: Deprecated parameter, has no effect.
+        :param urn: The urn string
+        """
+
+        # TODO: Deprecate the `type` parameter and remove it.
         if urn not in self._urns_set:
             self.urns.append(urn)
             self._urns_set.add(urn)
 
     def get_urns_not_in(
         self, type: str, other_checkpoint_state: "GenericCheckpointState"
     ) -> Iterable[str]:
+        """
+        Gets the urns present in this checkpoint but not the other_checkpoint for the given type.
+
+        :param type: Deprecated. Set to "*".
+        :param other_checkpoint_state: the checkpoint state to compute the urn set difference against.
+        :return: an iterable to the set of urns present in this checkpoint state but not in the other_checkpoint.
+        """
+
         diff = set(self.urns) - set(other_checkpoint_state.urns)
 
         # To maintain backwards compatibility, we provide this filtering mechanism.
+        # TODO: Deprecate the `type` parameter and remove it.
         if type == "*":
             yield from diff
         elif type == "topic":
@@ -110,6 +121,36 @@ def get_urns_not_in(
     def get_percent_entities_changed(
         self, old_checkpoint_state: "GenericCheckpointState"
     ) -> float:
-        return StaleEntityCheckpointStateBase.compute_percent_entities_changed(
+        """
+        Returns the percentage of entities that have changed relative to `old_checkpoint_state`.
+
+        :param old_checkpoint_state: the old checkpoint state to compute the relative change percent against.
+        :return: (1-|intersection(self, old_checkpoint_state)| / |old_checkpoint_state|) * 100.0
+        """
+        return compute_percent_entities_changed(
             [(self.urns, old_checkpoint_state.urns)]
         )
+
+
+def compute_percent_entities_changed(
+    new_old_entity_list: List[Tuple[List[str], List[str]]]
+) -> float:
+    old_count_all = 0
+    overlap_count_all = 0
+    for new_entities, old_entities in new_old_entity_list:
+        (overlap_count, old_count, _,) = get_entity_overlap_and_cardinalities(
+            new_entities=new_entities, old_entities=old_entities
+        )
+        overlap_count_all += overlap_count
+        old_count_all += old_count
+    if old_count_all:
+        return (1 - overlap_count_all / old_count_all) * 100.0
+    return 0.0
+
+
+def get_entity_overlap_and_cardinalities(
+    new_entities: List[str], old_entities: List[str]
+) -> Tuple[int, int, int]:
+    new_set = set(new_entities)
+    old_set = set(old_entities)
+    return len(new_set.intersection(old_set)), len(old_set), len(new_set)

metadata-ingestion/src/datahub/ingestion/source/state/profiling_state_handler.py

@@ -40,15 +40,17 @@ def __init__(
         pipeline_name: Optional[str],
         run_id: str,
     ):
-        self.source = source
+        self.state_provider = source.state_provider
         self.stateful_ingestion_config: Optional[
             ProfilingStatefulIngestionConfig
         ] = config.stateful_ingestion
         self.pipeline_name = pipeline_name
         self.run_id = run_id
-        self.checkpointing_enabled: bool = source.is_stateful_ingestion_configured()
+        self.checkpointing_enabled: bool = (
+            self.state_provider.is_stateful_ingestion_configured()
+        )
         self._job_id = self._init_job_id()
-        self.source.register_stateful_ingestion_usecase_handler(self)
+        self.state_provider.register_stateful_ingestion_usecase_handler(self)
 
     def _ignore_old_state(self) -> bool:
         if (
@@ -91,7 +93,7 @@ def create_checkpoint(self) -> Optional[Checkpoint[ProfilingCheckpointState]]:
     def get_current_state(self) -> Optional[ProfilingCheckpointState]:
         if not self.is_checkpointing_enabled() or self._ignore_new_state():
             return None
-        cur_checkpoint = self.source.get_current_checkpoint(self.job_id)
+        cur_checkpoint = self.state_provider.get_current_checkpoint(self.job_id)
         assert cur_checkpoint is not None
         cur_state = cast(ProfilingCheckpointState, cur_checkpoint.state)
         return cur_state
@@ -108,7 +110,7 @@ def add_to_state(
     def get_last_state(self) -> Optional[ProfilingCheckpointState]:
         if not self.is_checkpointing_enabled() or self._ignore_old_state():
             return None
-        last_checkpoint = self.source.get_last_checkpoint(
+        last_checkpoint = self.state_provider.get_last_checkpoint(
             self.job_id, ProfilingCheckpointState
         )
         if last_checkpoint and last_checkpoint.state:

metadata-ingestion/src/datahub/ingestion/source/state/redundant_run_skip_handler.py

@@ -46,14 +46,17 @@ def __init__(
         run_id: str,
     ):
         self.source = source
+        self.state_provider = source.state_provider
         self.stateful_ingestion_config: Optional[
             StatefulRedundantRunSkipConfig
         ] = config.stateful_ingestion
         self.pipeline_name = pipeline_name
         self.run_id = run_id
-        self.checkpointing_enabled: bool = source.is_stateful_ingestion_configured()
+        self.checkpointing_enabled: bool = (
+            self.state_provider.is_stateful_ingestion_configured()
+        )
         self._job_id = self._init_job_id()
-        self.source.register_stateful_ingestion_usecase_handler(self)
+        self.state_provider.register_stateful_ingestion_usecase_handler(self)
 
     def _ignore_old_state(self) -> bool:
         if (
@@ -114,7 +117,7 @@ def update_state(
     ) -> None:
         if not self.is_checkpointing_enabled() or self._ignore_new_state():
             return
-        cur_checkpoint = self.source.get_current_checkpoint(self.job_id)
+        cur_checkpoint = self.state_provider.get_current_checkpoint(self.job_id)
         assert cur_checkpoint is not None
         cur_state = cast(BaseUsageCheckpointState, cur_checkpoint.state)
         cur_state.begin_timestamp_millis = start_time_millis
@@ -125,7 +128,7 @@ def should_skip_this_run(self, cur_start_time_millis: int) -> bool:
             return False
         # Determine from the last check point state
         last_successful_pipeline_run_end_time_millis: Optional[int] = None
-        last_checkpoint = self.source.get_last_checkpoint(
+        last_checkpoint = self.state_provider.get_last_checkpoint(
             self.job_id, BaseUsageCheckpointState
         )
         if last_checkpoint and last_checkpoint.state: