acryl-datahub 1.0.0rc5__py3-none-any.whl → 1.0.0rc7__py3-none-any.whl

datahub/ingestion/source/dremio/dremio_api.py

@@ -271,12 +271,12 @@ class DremioAPIOperations:
                     self.cancel_query(job_id)
                     raise DremioAPIException(
                         f"Query execution timed out after {timeout} seconds"
-                    )
+                    ) from None
                 except RuntimeError as e:
-                    raise DremioAPIException(f"{str(e)}")
+                    raise DremioAPIException() from e
 
         except requests.RequestException as e:
-            raise DremioAPIException(f"Error executing query: {str(e)}")
+            raise DremioAPIException("Error executing query") from e
 
     def fetch_results(self, job_id: str) -> List[Dict]:
         """Fetch job results with status checking"""

datahub/ingestion/source/dremio/dremio_aspects.py

@@ -168,8 +168,9 @@ class DremioAspects:
         )
 
     def get_container_urn(
-        self, name: Optional[str] = None, path: Optional[List[str]] = []
+        self, name: Optional[str] = None, path: Optional[List[str]] = None
     ) -> str:
+        path = path or []
         container_key = self.get_container_key(name, path)
         return container_key.as_urn()
 

datahub/ingestion/source/elastic_search.py

@@ -32,9 +32,17 @@ 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,
+    StaleEntityRemovalSourceReport,
+)
+from datahub.ingestion.source.state.stateful_ingestion_base import (
+    StatefulIngestionConfigBase,
+    StatefulIngestionSourceBase,
+)
 from datahub.ingestion.source_config.operation_config import (
     OperationConfig,
     is_profiling_enabled,
@@ -188,7 +196,7 @@ class ElasticToSchemaFieldConverter:
 
 
 @dataclass
-class ElasticsearchSourceReport(SourceReport):
+class ElasticsearchSourceReport(StaleEntityRemovalSourceReport):
     index_scanned: int = 0
     filtered: LossyList[str] = field(default_factory=LossyList)
 
@@ -240,7 +248,11 @@ def collapse_urn(urn: str, collapse_urns: CollapseUrns) -> str:
     )
 
 
-class ElasticsearchSourceConfig(PlatformInstanceConfigMixin, EnvConfigMixin):
+class ElasticsearchSourceConfig(
+    StatefulIngestionConfigBase,
+    PlatformInstanceConfigMixin,
+    EnvConfigMixin,
+):
     host: str = Field(
         default="localhost:9200", description="The elastic search host URI."
     )
@@ -337,7 +349,7 @@ class ElasticsearchSourceConfig(PlatformInstanceConfigMixin, EnvConfigMixin):
 @config_class(ElasticsearchSourceConfig)
 @support_status(SupportStatus.CERTIFIED)
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
-class ElasticsearchSource(Source):
+class ElasticsearchSource(StatefulIngestionSourceBase):
     """
     This plugin extracts the following:
 
@@ -346,7 +358,7 @@ class ElasticsearchSource(Source):
     """
 
     def __init__(self, config: ElasticsearchSourceConfig, ctx: PipelineContext):
-        super().__init__(ctx)
+        super().__init__(config, ctx)
         self.source_config = config
         self.client = Elasticsearch(
             self.source_config.host,
@@ -361,7 +373,7 @@ class ElasticsearchSource(Source):
             ssl_assert_fingerprint=self.source_config.ssl_assert_fingerprint,
             url_prefix=self.source_config.url_prefix,
         )
-        self.report = ElasticsearchSourceReport()
+        self.report: ElasticsearchSourceReport = ElasticsearchSourceReport()
         self.data_stream_partition_count: Dict[str, int] = defaultdict(int)
         self.platform: str = "elasticsearch"
         self.cat_response: Optional[List[Dict[str, Any]]] = None
@@ -373,6 +385,14 @@ class ElasticsearchSource(Source):
         config = ElasticsearchSourceConfig.parse_obj(config_dict)
         return cls(config, ctx)
 
+    def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
+        return [
+            *super().get_workunit_processors(),
+            StaleEntityRemovalHandler.create(
+                self, self.source_config, self.ctx
+            ).workunit_processor,
+        ]
+
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
         indices = self.client.indices.get_alias()
         for index in indices:

datahub/ingestion/source/feast.py

@@ -20,7 +20,6 @@ from feast.data_source import DataSource
 from pydantic import Field
 
 import datahub.emitter.mce_builder as builder
-from datahub.configuration.common import ConfigModel
 from datahub.emitter.mce_builder import DEFAULT_ENV
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
@@ -31,8 +30,16 @@ 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, SourceReport
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+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.com.linkedin.pegasus2avro.common import MLFeatureDataType
 from datahub.metadata.com.linkedin.pegasus2avro.metadata.snapshot import (
     MLFeatureSnapshot,
@@ -86,7 +93,9 @@ _field_type_mapping: Dict[Union[ValueType, feast.types.FeastType], str] = {
 }
 
 
-class FeastRepositorySourceConfig(ConfigModel):
+class FeastRepositorySourceConfig(
+    StatefulIngestionConfigBase,
+):
     path: str = Field(description="Path to Feast repository")
     fs_yaml_file: Optional[str] = Field(
         default=None,
@@ -122,7 +131,7 @@ class FeastRepositorySourceConfig(ConfigModel):
 @capability(SourceCapability.SCHEMA_METADATA, "Enabled by default")
 @capability(SourceCapability.LINEAGE_COARSE, "Enabled by default")
 @dataclass
-class FeastRepositorySource(Source):
+class FeastRepositorySource(StatefulIngestionSourceBase):
     """
     This plugin extracts:
 
@@ -135,13 +144,14 @@ class FeastRepositorySource(Source):
 
     platform = "feast"
     source_config: FeastRepositorySourceConfig
-    report: SourceReport
+    report: StaleEntityRemovalSourceReport
     feature_store: FeatureStore
 
     def __init__(self, config: FeastRepositorySourceConfig, ctx: PipelineContext):
-        super().__init__(ctx)
+        super().__init__(config, ctx)
         self.source_config = config
-        self.report = SourceReport()
+        self.ctx = ctx
+        self.report = StaleEntityRemovalSourceReport()
         self.feature_store = FeatureStore(
             repo_path=self.source_config.path,
             fs_yaml_file=self.source_config.fs_yaml_file,
@@ -158,7 +168,8 @@ class FeastRepositorySource(Source):
 
         if ml_feature_data_type is None:
             self.report.report_warning(
-                parent_name, f"unable to map type {field_type} to metadata schema"
+                "unable to map type",
+                f"unable to map type {field_type} to metadata schema to parent: {parent_name}",
             )
 
             ml_feature_data_type = MLFeatureDataType.UNKNOWN
@@ -456,6 +467,14 @@ class FeastRepositorySource(Source):
         config = FeastRepositorySourceConfig.parse_obj(config_dict)
         return cls(config, ctx)
 
+    def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
+        return [
+            *super().get_workunit_processors(),
+            StaleEntityRemovalHandler.create(
+                self, self.source_config, self.ctx
+            ).workunit_processor,
+        ]
+
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
         for feature_view in self.feature_store.list_feature_views():
             for entity_name in feature_view.entities:

datahub/ingestion/source/file.py

@@ -351,7 +351,7 @@ class GenericFileSource(StatefulIngestionSourceBase, TestableSource):
                     self.report.add_deserialize_time(deserialize_duration)
                     yield i, item
             except Exception as e:
-                self.report.report_failure(f"path-{i}", str(e))
+                self.report.report_failure(f"{file_status.path}-{i}", str(e))
 
     @staticmethod
     def test_connection(config_dict: dict) -> TestConnectionReport:

datahub/ingestion/source/gc/execution_request_cleanup.py

@@ -130,8 +130,9 @@ class DatahubExecutionRequestCleanup:
         )
 
     def _scroll_execution_requests(
-        self, overrides: Dict[str, Any] = {}
+        self, overrides: Optional[Dict[str, Any]] = None
     ) -> Iterator[CleanupRecord]:
+        overrides = overrides or {}
         headers: Dict[str, Any] = {
             "Accept": "application/json",
             "Content-Type": "application/json",

datahub/ingestion/source/identity/okta.py

@@ -14,7 +14,6 @@ from okta.models import Group, GroupProfile, User, UserProfile, UserStatus
 from pydantic import validator
 from pydantic.fields import Field
 
-from datahub.configuration.common import ConfigModel
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
@@ -56,7 +55,7 @@ logger = logging.getLogger(__name__)
 nest_asyncio.apply()
 
 
-class OktaConfig(StatefulIngestionConfigBase, ConfigModel):
+class OktaConfig(StatefulIngestionConfigBase):
     # Required: Domain of the Okta deployment. Example: dev-33231928.okta.com
     okta_domain: str = Field(
         description="The location of your Okta Domain, without a protocol. Can be found in Okta Developer console. e.g. dev-33231928.okta.com",

datahub/ingestion/source/kafka/kafka.py

@@ -272,7 +272,7 @@ class KafkaSource(StatefulIngestionSourceBase, TestableSource):
             return schema_registry_class.create(config, report)
         except Exception as e:
             logger.debug(e, exc_info=e)
-            raise ImportError(config.schema_registry_class)
+            raise ImportError(config.schema_registry_class) from e
 
     def __init__(self, config: KafkaSourceConfig, ctx: PipelineContext):
         super().__init__(config, ctx)

datahub/ingestion/source/looker/looker_file_loader.py

@@ -33,14 +33,14 @@ class LookerViewFileLoader:
         base_projects_folder: Dict[str, pathlib.Path],
         reporter: LookMLSourceReport,
         source_config: LookMLSourceConfig,
-        manifest_constants: Dict[str, LookerConstant] = {},
+        manifest_constants: Optional[Dict[str, LookerConstant]] = None,
     ) -> None:
         self.viewfile_cache: Dict[str, Optional[LookerViewFile]] = {}
         self._root_project_name = root_project_name
         self._base_projects_folder = base_projects_folder
         self.reporter = reporter
         self.source_config = source_config
-        self.manifest_constants = manifest_constants
+        self.manifest_constants = manifest_constants or {}
 
     def _load_viewfile(
         self, project_name: str, path: str, reporter: LookMLSourceReport

datahub/ingestion/source/looker/lookml_source.py

@@ -501,7 +501,7 @@ class LookMLSource(StatefulIngestionSourceBase):
             raise ValueError(
                 f"Could not locate a project name for model {model_name}. Consider configuring a static project name "
                 f"in your config file"
-            )
+            ) from None
 
     def get_manifest_if_present(self, folder: pathlib.Path) -> Optional[LookerManifest]:
         manifest_file = folder / "manifest.lkml"

datahub/ingestion/source/metabase.py

@@ -69,9 +69,19 @@ class MetabaseConfig(DatasetLineageProviderConfigBase, StatefulIngestionConfigBa
         default=None,
         description="optional URL to use in links (if `connect_uri` is only for ingestion)",
     )
-    username: Optional[str] = Field(default=None, description="Metabase username.")
+    username: Optional[str] = Field(
+        default=None,
+        description="Metabase username, used when an API key is not provided.",
+    )
     password: Optional[pydantic.SecretStr] = Field(
-        default=None, description="Metabase password."
+        default=None,
+        description="Metabase password, used when an API key is not provided.",
+    )
+
+    # https://www.metabase.com/learn/metabase-basics/administration/administration-and-operation/metabase-api#example-get-request
+    api_key: Optional[pydantic.SecretStr] = Field(
+        default=None,
+        description="Metabase API key. If provided, the username and password will be ignored. Recommended method.",
     )
     # TODO: Check and remove this if no longer needed.
     # Config database_alias is removed from sql sources.
@@ -178,30 +188,40 @@ class MetabaseSource(StatefulIngestionSourceBase):
         self.source_config: MetabaseConfig = config
 
     def setup_session(self) -> None:
-        login_response = requests.post(
-            f"{self.config.connect_uri}/api/session",
-            None,
-            {
-                "username": self.config.username,
-                "password": (
-                    self.config.password.get_secret_value()
-                    if self.config.password
-                    else None
-                ),
-            },
-        )
+        self.session = requests.session()
+        if self.config.api_key:
+            self.session.headers.update(
+                {
+                    "x-api-key": self.config.api_key.get_secret_value(),
+                    "Content-Type": "application/json",
+                    "Accept": "*/*",
+                }
+            )
+        else:
+            # If no API key is provided, generate a session token using username and password.
+            login_response = requests.post(
+                f"{self.config.connect_uri}/api/session",
+                None,
+                {
+                    "username": self.config.username,
+                    "password": (
+                        self.config.password.get_secret_value()
+                        if self.config.password
+                        else None
+                    ),
+                },
+            )
 
-        login_response.raise_for_status()
-        self.access_token = login_response.json().get("id", "")
+            login_response.raise_for_status()
+            self.access_token = login_response.json().get("id", "")
 
-        self.session = requests.session()
-        self.session.headers.update(
-            {
-                "X-Metabase-Session": f"{self.access_token}",
-                "Content-Type": "application/json",
-                "Accept": "*/*",
-            }
-        )
+            self.session.headers.update(
+                {
+                    "X-Metabase-Session": f"{self.access_token}",
+                    "Content-Type": "application/json",
+                    "Accept": "*/*",
+                }
+            )
 
         # Test the connection
         try:
@@ -217,15 +237,17 @@ class MetabaseSource(StatefulIngestionSourceBase):
             )
 
     def close(self) -> None:
-        response = requests.delete(
-            f"{self.config.connect_uri}/api/session",
-            headers={"X-Metabase-Session": self.access_token},
-        )
-        if response.status_code not in (200, 204):
-            self.report.report_failure(
-                title="Unable to Log User Out",
-                message=f"Unable to logout for user {self.config.username}",
+        # API key authentication does not require session closure.
+        if not self.config.api_key:
+            response = requests.delete(
+                f"{self.config.connect_uri}/api/session",
+                headers={"X-Metabase-Session": self.access_token},
             )
+            if response.status_code not in (200, 204):
+                self.report.report_failure(
+                    title="Unable to Log User Out",
+                    message=f"Unable to logout for user {self.config.username}",
+                )
         super().close()
 
     def emit_dashboard_mces(self) -> Iterable[MetadataWorkUnit]:

datahub/ingestion/source/mlflow.py

@@ -1,5 +1,5 @@
 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
@@ -8,7 +8,9 @@ from mlflow.store.entities import PagedList
 from pydantic.fields import Field
 
 import datahub.emitter.mce_builder as builder
-from datahub.configuration.source_common import EnvConfigMixin
+from datahub.configuration.source_common import (
+    EnvConfigMixin,
+)
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
@@ -18,8 +20,20 @@ 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.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 (
     GlobalTagsClass,
     MLHyperParamClass,
@@ -35,7 +49,7 @@ from datahub.metadata.schema_classes import (
 T = TypeVar("T")
 
 
-class MLflowConfig(EnvConfigMixin):
+class MLflowConfig(StatefulIngestionConfigBase, EnvConfigMixin):
     tracking_uri: Optional[str] = Field(
         default=None,
         description=(
@@ -79,7 +93,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 +119,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,6 +131,14 @@ 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_ml_model_workunits()

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(
@@ -1489,7 +1494,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):
@@ -282,6 +294,14 @@ class Neo4jSource(Source):
     def get_relationships(self, record: dict) -> dict:
         return record.get("relationships", None)
 
+    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(
             "CALL apoc.meta.schema() YIELD value UNWIND keys(value) AS key RETURN key, value[key] AS value;"

datahub/ingestion/source/nifi.py

@@ -22,7 +22,9 @@ from requests_gssapi import HTTPSPNEGOAuth
 
 import datahub.emitter.mce_builder as builder
 from datahub.configuration.common import AllowDenyPattern
-from datahub.configuration.source_common import EnvConfigMixin
+from datahub.configuration.source_common import (
+    EnvConfigMixin,
+)
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_builder import ContainerKey, gen_containers
 from datahub.ingestion.api.common import PipelineContext
@@ -33,9 +35,21 @@ 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 JobContainerSubTypes
+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 (
     BrowsePathEntryClass,
     BrowsePathsV2Class,
@@ -81,7 +95,7 @@ class ProcessGroupKey(ContainerKey):
     process_group_id: str
 
 
-class NifiSourceConfig(EnvConfigMixin):
+class NifiSourceConfig(StatefulIngestionConfigBase, EnvConfigMixin):
     site_url: str = Field(
         description="URL for Nifi, ending with /nifi/. e.g. https://mynifi.domain/nifi/"
     )
@@ -452,7 +466,7 @@ def get_attribute_value(attr_lst: List[dict], attr_name: str) -> Optional[str]:
 
 
 @dataclass
-class NifiSourceReport(SourceReport):
+class NifiSourceReport(StaleEntityRemovalSourceReport):
     filtered: LossyList[str] = field(default_factory=LossyList)
 
     def report_dropped(self, ent_name: str) -> None:
@@ -464,13 +478,14 @@ class NifiSourceReport(SourceReport):
 @config_class(NifiSourceConfig)
 @support_status(SupportStatus.CERTIFIED)
 @capability(SourceCapability.LINEAGE_COARSE, "Supported. See docs for limitations")
-class NifiSource(Source):
+class NifiSource(StatefulIngestionSourceBase):
     config: NifiSourceConfig
     report: NifiSourceReport
 
     def __init__(self, config: NifiSourceConfig, ctx: PipelineContext) -> None:
-        super().__init__(ctx)
+        super().__init__(config, ctx)
         self.config = config
+        self.ctx = ctx
         self.report = NifiSourceReport()
         self.session = requests.Session()
 
@@ -1151,6 +1166,14 @@ class NifiSource(Source):
         token_response.raise_for_status()
         self.session.headers.update({"Authorization": "Bearer " + token_response.text})
 
+    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]:
         try:
             self.authenticate()