acryl-datahub 0.15.0.6rc2__py3-none-any.whl → 1.0.0__py3-none-any.whl

datahub/ingestion/source/mlflow.py

@@ -1,15 +1,20 @@
+import time
 from dataclasses import dataclass
-from typing import Any, Callable, Iterable, Optional, TypeVar, Union
+from typing import Any, Callable, Iterable, List, Optional, TypeVar, Union
 
 from mlflow import MlflowClient
-from mlflow.entities import Run
+from mlflow.entities import Experiment, Run
 from mlflow.entities.model_registry import ModelVersion, RegisteredModel
 from mlflow.store.entities import PagedList
 from pydantic.fields import Field
 
 import datahub.emitter.mce_builder as builder
+from datahub.api.entities.dataprocess.dataprocess_instance import (
+    DataProcessInstance,
+)
 from datahub.configuration.source_common import EnvConfigMixin
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.emitter.mcp_builder import ContainerKey
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SupportStatus,
@@ -18,24 +23,62 @@ from datahub.ingestion.api.decorators import (
     platform_name,
     support_status,
 )
-from datahub.ingestion.api.source import Source, SourceCapability, SourceReport
+from datahub.ingestion.api.source import (
+    MetadataWorkUnitProcessor,
+    SourceCapability,
+    SourceReport,
+)
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import MLAssetSubTypes
+from datahub.ingestion.source.state.stale_entity_removal_handler import (
+    StaleEntityRemovalHandler,
+    StaleEntityRemovalSourceReport,
+)
+from datahub.ingestion.source.state.stateful_ingestion_base import (
+    StatefulIngestionConfigBase,
+    StatefulIngestionSourceBase,
+)
 from datahub.metadata.schema_classes import (
+    AuditStampClass,
+    ContainerClass,
+    DataPlatformInstanceClass,
+    DataProcessInstanceOutputClass,
+    DataProcessInstancePropertiesClass,
+    DataProcessInstanceRunEventClass,
+    DataProcessInstanceRunResultClass,
+    DataProcessRunStatusClass,
+    EdgeClass,
     GlobalTagsClass,
+    MetadataAttributionClass,
     MLHyperParamClass,
     MLMetricClass,
     MLModelGroupPropertiesClass,
     MLModelPropertiesClass,
+    MLTrainingRunPropertiesClass,
+    PlatformResourceInfoClass,
+    SubTypesClass,
     TagAssociationClass,
     TagPropertiesClass,
+    TimeStampClass,
+    VersionPropertiesClass,
     VersionTagClass,
     _Aspect,
 )
+from datahub.metadata.urns import (
+    DataPlatformUrn,
+    MlModelUrn,
+    VersionSetUrn,
+)
+from datahub.sdk.container import Container
 
 T = TypeVar("T")
 
 
-class MLflowConfig(EnvConfigMixin):
+class ContainerKeyWithId(ContainerKey):
+    id: str
+
+
+class MLflowConfig(StatefulIngestionConfigBase, EnvConfigMixin):
     tracking_uri: Optional[str] = Field(
         default=None,
         description=(
@@ -79,7 +122,7 @@ class MLflowRegisteredModelStageInfo:
     "Extract descriptions for MLflow Registered Models and Model Versions",
 )
 @capability(SourceCapability.TAGS, "Extract tags for MLflow Registered Model Stages")
-class MLflowSource(Source):
+class MLflowSource(StatefulIngestionSourceBase):
     platform = "mlflow"
     registered_model_stages_info = (
         MLflowRegisteredModelStageInfo(
@@ -105,9 +148,10 @@ class MLflowSource(Source):
     )
 
     def __init__(self, ctx: PipelineContext, config: MLflowConfig):
-        super().__init__(ctx)
+        super().__init__(config, ctx)
+        self.ctx = ctx
         self.config = config
-        self.report = SourceReport()
+        self.report = StaleEntityRemovalSourceReport()
         self.client = MlflowClient(
             tracking_uri=self.config.tracking_uri,
             registry_uri=self.config.registry_uri,
@@ -116,8 +160,17 @@ class MLflowSource(Source):
     def get_report(self) -> SourceReport:
         return self.report
 
+    def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
+        return [
+            *super().get_workunit_processors(),
+            StaleEntityRemovalHandler.create(
+                self, self.config, self.ctx
+            ).workunit_processor,
+        ]
+
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
         yield from self._get_tags_workunits()
+        yield from self._get_experiment_workunits()
         yield from self._get_ml_model_workunits()
 
     def _get_tags_workunits(self) -> Iterable[MetadataWorkUnit]:
@@ -151,22 +204,162 @@ class MLflowSource(Source):
             aspect=aspect,
         ).as_workunit()
 
-    def _get_ml_model_workunits(self) -> Iterable[MetadataWorkUnit]:
-        """
-        Traverse each Registered Model in Model Registry and generate a corresponding workunit.
-        """
-        registered_models = self._get_mlflow_registered_models()
-        for registered_model in registered_models:
-            yield self._get_ml_group_workunit(registered_model)
-            model_versions = self._get_mlflow_model_versions(registered_model)
-            for model_version in model_versions:
-                run = self._get_mlflow_run(model_version)
-                yield self._get_ml_model_properties_workunit(
-                    registered_model=registered_model,
-                    model_version=model_version,
-                    run=run,
-                )
-                yield self._get_global_tags_workunit(model_version=model_version)
+    def _get_experiment_workunits(self) -> Iterable[MetadataWorkUnit]:
+        experiments = self._get_mlflow_experiments()
+        for experiment in experiments:
+            yield from self._get_experiment_container_workunit(experiment)
+
+            runs = self._get_mlflow_runs_from_experiment(experiment)
+            if runs:
+                for run in runs:
+                    yield from self._get_run_workunits(experiment, run)
+
+    def _get_experiment_custom_properties(self, experiment):
+        experiment_custom_props = getattr(experiment, "tags", {}) or {}
+        experiment_custom_props.pop("mlflow.note.content", None)
+        experiment_custom_props["artifacts_location"] = experiment.artifact_location
+        return experiment_custom_props
+
+    def _get_experiment_container_workunit(
+        self, experiment: Experiment
+    ) -> Iterable[MetadataWorkUnit]:
+        experiment_container = Container(
+            container_key=ContainerKeyWithId(
+                platform=str(DataPlatformUrn(platform_name=self.platform)),
+                id=experiment.name,
+            ),
+            subtype=MLAssetSubTypes.MLFLOW_EXPERIMENT,
+            display_name=experiment.name,
+            description=experiment.tags.get("mlflow.note.content"),
+            extra_properties=self._get_experiment_custom_properties(experiment),
+        )
+
+        yield from experiment_container.as_workunits()
+
+    def _get_run_metrics(self, run: Run) -> List[MLMetricClass]:
+        return [
+            MLMetricClass(name=k, value=str(v)) for k, v in run.data.metrics.items()
+        ]
+
+    def _get_run_params(self, run: Run) -> List[MLHyperParamClass]:
+        return [
+            MLHyperParamClass(name=k, value=str(v)) for k, v in run.data.params.items()
+        ]
+
+    def _convert_run_result_type(
+        self, status: str
+    ) -> DataProcessInstanceRunResultClass:
+        if status == "FINISHED":
+            return DataProcessInstanceRunResultClass(
+                type="SUCCESS", nativeResultType=self.platform
+            )
+        elif status == "FAILED":
+            return DataProcessInstanceRunResultClass(
+                type="FAILURE", nativeResultType=self.platform
+            )
+        else:
+            return DataProcessInstanceRunResultClass(
+                type="SKIPPED", nativeResultType=self.platform
+            )
+
+    def _get_run_workunits(
+        self, experiment: Experiment, run: Run
+    ) -> Iterable[MetadataWorkUnit]:
+        experiment_key = ContainerKeyWithId(
+            platform=str(DataPlatformUrn(self.platform)), id=experiment.name
+        )
+
+        data_process_instance = DataProcessInstance(
+            id=run.info.run_id,
+            orchestrator=self.platform,
+            template_urn=None,
+        )
+
+        created_time = run.info.start_time or int(time.time() * 1000)
+        user_id = run.info.user_id if run.info.user_id else "mlflow"
+        guid_dict_user = {"platform": self.platform, "user": user_id}
+        platform_user_urn = (
+            f"urn:li:platformResource:{builder.datahub_guid(guid_dict_user)}"
+        )
+
+        yield MetadataChangeProposalWrapper(
+            entityUrn=platform_user_urn,
+            aspect=PlatformResourceInfoClass(
+                resourceType="user",
+                primaryKey=user_id,
+            ),
+        ).as_workunit()
+
+        yield MetadataChangeProposalWrapper(
+            entityUrn=str(data_process_instance.urn),
+            aspect=DataProcessInstancePropertiesClass(
+                name=run.info.run_name or run.info.run_id,
+                created=AuditStampClass(
+                    time=created_time,
+                    actor=platform_user_urn,
+                ),
+                externalUrl=self._make_external_url_from_run(experiment, run),
+                customProperties=getattr(run, "tags", {}) or {},
+            ),
+        ).as_workunit()
+
+        yield MetadataChangeProposalWrapper(
+            entityUrn=str(data_process_instance.urn),
+            aspect=ContainerClass(container=experiment_key.as_urn()),
+        ).as_workunit()
+
+        model_versions = self.get_mlflow_model_versions_from_run(run.info.run_id)
+        if model_versions:
+            model_version_urn = self._make_ml_model_urn(model_versions[0])
+            yield MetadataChangeProposalWrapper(
+                entityUrn=str(data_process_instance.urn),
+                aspect=DataProcessInstanceOutputClass(
+                    outputs=[],
+                    outputEdges=[
+                        EdgeClass(destinationUrn=model_version_urn),
+                    ],
+                ),
+            ).as_workunit()
+
+        metrics = self._get_run_metrics(run)
+        hyperparams = self._get_run_params(run)
+        yield MetadataChangeProposalWrapper(
+            entityUrn=str(data_process_instance.urn),
+            aspect=MLTrainingRunPropertiesClass(
+                hyperParams=hyperparams,
+                trainingMetrics=metrics,
+                outputUrls=[run.info.artifact_uri],
+                id=run.info.run_id,
+            ),
+        ).as_workunit()
+
+        if run.info.end_time:
+            duration_millis = run.info.end_time - run.info.start_time
+
+            yield MetadataChangeProposalWrapper(
+                entityUrn=str(data_process_instance.urn),
+                aspect=DataProcessInstanceRunEventClass(
+                    status=DataProcessRunStatusClass.COMPLETE,
+                    timestampMillis=run.info.end_time,
+                    result=DataProcessInstanceRunResultClass(
+                        type=self._convert_run_result_type(run.info.status).type,
+                        nativeResultType=self.platform,
+                    ),
+                    durationMillis=duration_millis,
+                ),
+            ).as_workunit()
+
+        yield MetadataChangeProposalWrapper(
+            entityUrn=str(data_process_instance.urn),
+            aspect=DataPlatformInstanceClass(
+                platform=str(DataPlatformUrn(self.platform))
+            ),
+        ).as_workunit()
+
+        yield MetadataChangeProposalWrapper(
+            entityUrn=str(data_process_instance.urn),
+            aspect=SubTypesClass(typeNames=[MLAssetSubTypes.MLFLOW_TRAINING_RUN]),
+        ).as_workunit()
 
     def _get_mlflow_registered_models(self) -> Iterable[RegisteredModel]:
         """
@@ -179,6 +372,19 @@ class MLflowSource(Source):
         )
         return registered_models
 
+    def _get_mlflow_experiments(self) -> Iterable[Experiment]:
+        experiments: Iterable[Experiment] = self._traverse_mlflow_search_func(
+            search_func=self.client.search_experiments,
+        )
+        return experiments
+
+    def _get_mlflow_runs_from_experiment(self, experiment: Experiment) -> Iterable[Run]:
+        runs: Iterable[Run] = self._traverse_mlflow_search_func(
+            search_func=self.client.search_runs,
+            experiment_ids=[experiment.experiment_id],
+        )
+        return runs
+
     @staticmethod
     def _traverse_mlflow_search_func(
         search_func: Callable[..., PagedList[T]],
@@ -195,6 +401,13 @@ class MLflowSource(Source):
             if not next_page_token:
                 return
 
+    def _get_latest_version(self, registered_model: RegisteredModel) -> Optional[str]:
+        return (
+            str(registered_model.latest_versions[0].version)
+            if registered_model.latest_versions
+            else None
+        )
+
     def _get_ml_group_workunit(
         self,
         registered_model: RegisteredModel,
@@ -206,7 +419,20 @@ class MLflowSource(Source):
         ml_model_group_properties = MLModelGroupPropertiesClass(
             customProperties=registered_model.tags,
             description=registered_model.description,
-            createdAt=registered_model.creation_timestamp,
+            created=TimeStampClass(
+                time=registered_model.creation_timestamp, actor=None
+            ),
+            lastModified=TimeStampClass(
+                time=registered_model.last_updated_timestamp,
+                actor=None,
+            ),
+            version=VersionTagClass(
+                versionTag=self._get_latest_version(registered_model),
+                metadataAttribution=MetadataAttributionClass(
+                    time=registered_model.last_updated_timestamp,
+                    actor="urn:li:corpuser:datahub",
+                ),
+            ),
         )
         wu = self._create_workunit(
             urn=ml_model_group_urn,
@@ -236,6 +462,16 @@ class MLflowSource(Source):
         )
         return model_versions
 
+    def get_mlflow_model_versions_from_run(self, run_id):
+        filter_string = f"run_id = '{run_id}'"
+
+        model_versions: Iterable[ModelVersion] = self._traverse_mlflow_search_func(
+            search_func=self.client.search_model_versions,
+            filter_string=filter_string,
+        )
+
+        return list(model_versions)
+
     def _get_mlflow_run(self, model_version: ModelVersion) -> Union[None, Run]:
         """
         Get a Run associated with a Model Version. Some MVs may exist without Run.
@@ -246,6 +482,67 @@ class MLflowSource(Source):
         else:
             return None
 
+    def _get_ml_model_workunits(self) -> Iterable[MetadataWorkUnit]:
+        """
+        Traverse each Registered Model in Model Registry and generate a corresponding workunit.
+        """
+        registered_models = self._get_mlflow_registered_models()
+        for registered_model in registered_models:
+            version_set_urn = self._get_version_set_urn(registered_model)
+            yield self._get_ml_group_workunit(registered_model)
+            model_versions = self._get_mlflow_model_versions(registered_model)
+            for model_version in model_versions:
+                run = self._get_mlflow_run(model_version)
+                yield self._get_ml_model_properties_workunit(
+                    registered_model=registered_model,
+                    model_version=model_version,
+                    run=run,
+                )
+                yield self._get_ml_model_version_properties_workunit(
+                    model_version=model_version,
+                    version_set_urn=version_set_urn,
+                )
+                yield self._get_global_tags_workunit(model_version=model_version)
+
+    def _get_version_set_urn(self, registered_model: RegisteredModel) -> VersionSetUrn:
+        guid_dict = {"platform": self.platform, "name": registered_model.name}
+        version_set_urn = VersionSetUrn(
+            id=builder.datahub_guid(guid_dict),
+            entity_type=MlModelUrn.ENTITY_TYPE,
+        )
+
+        return version_set_urn
+
+    def _get_ml_model_version_properties_workunit(
+        self,
+        model_version: ModelVersion,
+        version_set_urn: VersionSetUrn,
+    ) -> MetadataWorkUnit:
+        ml_model_urn = self._make_ml_model_urn(model_version)
+
+        # get mlmodel name from ml model urn
+        ml_model_version_properties = VersionPropertiesClass(
+            version=VersionTagClass(
+                versionTag=str(model_version.version),
+                metadataAttribution=MetadataAttributionClass(
+                    time=model_version.creation_timestamp,
+                    actor="urn:li:corpuser:datahub",
+                ),
+            ),
+            versionSet=str(version_set_urn),
+            sortId=str(model_version.version).zfill(10),
+            aliases=[
+                VersionTagClass(versionTag=alias) for alias in model_version.aliases
+            ],
+        )
+
+        wu = MetadataChangeProposalWrapper(
+            entityUrn=str(ml_model_urn),
+            aspect=ml_model_version_properties,
+        ).as_workunit()
+
+        return wu
+
     def _get_ml_model_properties_workunit(
         self,
         registered_model: RegisteredModel,
@@ -259,28 +556,47 @@ class MLflowSource(Source):
         """
         ml_model_group_urn = self._make_ml_model_group_urn(registered_model)
         ml_model_urn = self._make_ml_model_urn(model_version)
+
         if run:
-            hyperparams = [
-                MLHyperParamClass(name=k, value=str(v))
-                for k, v in run.data.params.items()
-            ]
-            training_metrics = [
-                MLMetricClass(name=k, value=str(v)) for k, v in run.data.metrics.items()
-            ]
+            # Use the same metrics and hyperparams from the run
+            hyperparams = self._get_run_params(run)
+            training_metrics = self._get_run_metrics(run)
+            run_urn = DataProcessInstance(
+                id=run.info.run_id,
+                orchestrator=self.platform,
+            ).urn
+
+            training_jobs = [str(run_urn)] if run_urn else []
         else:
             hyperparams = None
             training_metrics = None
+            training_jobs = []
+
+        created_time = model_version.creation_timestamp
+        created_actor = (
+            f"urn:li:platformResource:{model_version.user_id}"
+            if model_version.user_id
+            else None
+        )
+        model_version_tags = [f"{k}:{v}" for k, v in model_version.tags.items()]
+
         ml_model_properties = MLModelPropertiesClass(
             customProperties=model_version.tags,
             externalUrl=self._make_external_url(model_version),
+            lastModified=TimeStampClass(
+                time=model_version.last_updated_timestamp,
+                actor=None,
+            ),
             description=model_version.description,
-            date=model_version.creation_timestamp,
-            version=VersionTagClass(versionTag=str(model_version.version)),
+            created=TimeStampClass(
+                time=created_time,
+                actor=created_actor,
+            ),
             hyperParams=hyperparams,
             trainingMetrics=training_metrics,
-            # mlflow tags are dicts, but datahub tags are lists. currently use only keys from mlflow tags
-            tags=list(model_version.tags.keys()),
+            tags=model_version_tags,
             groups=[ml_model_group_urn],
+            trainingJobs=training_jobs,
         )
         wu = self._create_workunit(urn=ml_model_urn, aspect=ml_model_properties)
         return wu
@@ -314,6 +630,15 @@ class MLflowSource(Source):
         else:
             return None
 
+    def _make_external_url_from_run(
+        self, experiment: Experiment, run: Run
+    ) -> Union[None, str]:
+        base_uri = self.client.tracking_uri
+        if base_uri.startswith("http"):
+            return f"{base_uri.rstrip('/')}/#/experiments/{experiment.experiment_id}/runs/{run.info.run_id}"
+        else:
+            return None
+
     def _get_global_tags_workunit(
         self,
         model_version: ModelVersion,
@@ -333,3 +658,8 @@ class MLflowSource(Source):
             aspect=global_tags,
         )
         return wu
+
+    @classmethod
+    def create(cls, config_dict: dict, ctx: PipelineContext) -> "MLflowSource":
+        config = MLflowConfig.parse_obj(config_dict)
+        return cls(ctx, config)

datahub/ingestion/source/mode.py

@@ -23,7 +23,9 @@ from tenacity import retry_if_exception_type, stop_after_attempt, wait_exponenti
 
 import datahub.emitter.mce_builder as builder
 from datahub.configuration.common import AllowDenyPattern, ConfigModel
-from datahub.configuration.source_common import DatasetLineageProviderConfigBase
+from datahub.configuration.source_common import (
+    DatasetLineageProviderConfigBase,
+)
 from datahub.configuration.validate_field_removal import pydantic_removed_field
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_builder import (
@@ -137,7 +139,10 @@ class ModeAPIConfig(ConfigModel):
     )
 
 
-class ModeConfig(StatefulIngestionConfigBase, DatasetLineageProviderConfigBase):
+class ModeConfig(
+    StatefulIngestionConfigBase,
+    DatasetLineageProviderConfigBase,
+):
     # See https://mode.com/developer/api-reference/authentication/
     # for authentication
     connect_uri: str = Field(
@@ -154,7 +159,12 @@ class ModeConfig(StatefulIngestionConfigBase, DatasetLineageProviderConfigBase):
     )
 
     workspace: str = Field(
-        description="The Mode workspace name. Find it in Settings > Workspace > Details."
+        description="The Mode workspace username. If you navigate to Workspace Settings > Details, "
+        "the url will be `https://app.mode.com/organizations/<workspace-username>`. "
+        # The lowercase comment is derived from a comment in a Mode API example.
+        # https://mode.com/developer/api-cookbook/management/get-all-reports/
+        # > "Note: workspace_name value should be all lowercase"
+        "This is distinct from the workspace's display name, and should be all lowercase."
     )
     _default_schema = pydantic_removed_field("default_schema")
 
@@ -372,7 +382,7 @@ class ModeSource(StatefulIngestionSourceBase):
         ]
 
     def _dashboard_urn(self, report_info: dict) -> str:
-        return builder.make_dashboard_urn(self.platform, report_info.get("id", ""))
+        return builder.make_dashboard_urn(self.platform, str(report_info.get("id", "")))
 
     def _parse_last_run_at(self, report_info: dict) -> Optional[int]:
         # Mode queries are refreshed, and that timestamp is reflected correctly here.
@@ -759,9 +769,9 @@ class ModeSource(StatefulIngestionSourceBase):
                 return platform, database
         else:
             self.report.report_warning(
-                title="Failed to create Data Platform Urn",
-                message=f"Cannot create datasource urn for datasource id: "
-                f"{data_source_id}",
+                title="Unable to construct upstream lineage",
+                message="We did not find a data source / connection with a matching ID, meaning that we do not know the platform/database to use in lineage.",
+                context=f"Data Source ID: {data_source_id}",
             )
         return None, None
 
@@ -1489,7 +1499,7 @@ class ModeSource(StatefulIngestionSourceBase):
                     sleep_time = error_response.headers.get("retry-after")
                     if sleep_time is not None:
                         time.sleep(float(sleep_time))
-                    raise HTTPError429
+                    raise HTTPError429 from None
 
                 raise http_error
 

datahub/ingestion/source/neo4j/neo4j_source.py

@@ -7,7 +7,9 @@ import pandas as pd
 from neo4j import GraphDatabase
 from pydantic.fields import Field
 
-from datahub.configuration.source_common import EnvConfigMixin
+from datahub.configuration.source_common import (
+    EnvConfigMixin,
+)
 from datahub.emitter.mce_builder import make_data_platform_urn, make_dataset_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
@@ -17,9 +19,19 @@ from datahub.ingestion.api.decorators import (
     platform_name,
     support_status,
 )
-from datahub.ingestion.api.source import Source, SourceReport
+from datahub.ingestion.api.source import (
+    MetadataWorkUnitProcessor,
+)
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.common.subtypes import DatasetSubTypes
+from datahub.ingestion.source.state.stale_entity_removal_handler import (
+    StaleEntityRemovalHandler,
+)
+from datahub.ingestion.source.state.stateful_ingestion_base import (
+    StatefulIngestionConfigBase,
+    StatefulIngestionReport,
+    StatefulIngestionSourceBase,
+)
 from datahub.metadata.com.linkedin.pegasus2avro.schema import SchemaFieldDataType
 from datahub.metadata.schema_classes import (
     AuditStampClass,
@@ -52,7 +64,7 @@ _type_mapping: Dict[Union[Type, str], Type] = {
 }
 
 
-class Neo4jConfig(EnvConfigMixin):
+class Neo4jConfig(EnvConfigMixin, StatefulIngestionConfigBase):
     username: str = Field(description="Neo4j Username")
     password: str = Field(description="Neo4j Password")
     uri: str = Field(description="The URI for the Neo4j server")
@@ -60,7 +72,7 @@ class Neo4jConfig(EnvConfigMixin):
 
 
 @dataclass
-class Neo4jSourceReport(SourceReport):
+class Neo4jSourceReport(StatefulIngestionReport):
     obj_failures: int = 0
     obj_created: int = 0
 
@@ -68,7 +80,7 @@ class Neo4jSourceReport(SourceReport):
 @platform_name("Neo4j", id="neo4j")
 @config_class(Neo4jConfig)
 @support_status(SupportStatus.CERTIFIED)
-class Neo4jSource(Source):
+class Neo4jSource(StatefulIngestionSourceBase):
     NODE = "node"
     RELATIONSHIP = "relationship"
     PLATFORM = "neo4j"
@@ -76,7 +88,7 @@ class Neo4jSource(Source):
     def __init__(self, ctx: PipelineContext, config: Neo4jConfig):
         self.ctx = ctx
         self.config = config
-        self.report = Neo4jSourceReport()
+        self.report: Neo4jSourceReport = Neo4jSourceReport()
 
     @classmethod
     def create(cls, config_dict, ctx):
@@ -280,7 +292,15 @@ class Neo4jSource(Source):
         return record["properties"]
 
     def get_relationships(self, record: dict) -> dict:
-        return record.get("relationships", None)
+        return record.get("relationships", {})
+
+    def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
+        return [
+            *super().get_workunit_processors(),
+            StaleEntityRemovalHandler.create(
+                self, self.config, self.ctx
+            ).workunit_processor,
+        ]
 
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
         df = self.get_neo4j_metadata(