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

datahub/ingestion/source/cassandra/cassandra.py

@@ -1,19 +1,14 @@
 import dataclasses
 import json
 import logging
-from typing import Any, Dict, Iterable, List, Optional
+from typing import Any, Dict, Iterable, List, Optional, Union
 
 from datahub.emitter.mce_builder import (
-    make_data_platform_urn,
-    make_dataplatform_instance_urn,
     make_dataset_urn_with_platform_instance,
     make_schema_field_urn,
 )
-from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_builder import (
     ContainerKey,
-    add_dataset_to_container,
-    gen_containers,
 )
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
@@ -31,6 +26,7 @@ from datahub.ingestion.source.cassandra.cassandra_api import (
     CassandraColumn,
     CassandraEntities,
     CassandraKeyspace,
+    CassandraSharedDatasetFields,
     CassandraTable,
     CassandraView,
 )
@@ -51,24 +47,21 @@ from datahub.ingestion.source.state.stale_entity_removal_handler import (
 from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionSourceBase,
 )
-from datahub.metadata.com.linkedin.pegasus2avro.common import StatusClass
 from datahub.metadata.com.linkedin.pegasus2avro.schema import (
     SchemaField,
-    SchemaMetadata,
 )
 from datahub.metadata.schema_classes import (
-    DataPlatformInstanceClass,
     DatasetLineageTypeClass,
-    DatasetPropertiesClass,
     FineGrainedLineageClass,
     FineGrainedLineageDownstreamTypeClass,
     FineGrainedLineageUpstreamTypeClass,
-    OtherSchemaClass,
-    SubTypesClass,
     UpstreamClass,
     UpstreamLineageClass,
     ViewPropertiesClass,
 )
+from datahub.sdk._entity import Entity
+from datahub.sdk.container import Container
+from datahub.sdk.dataset import Dataset
 
 logger = logging.getLogger(__name__)
 
@@ -133,6 +126,13 @@ class CassandraSource(StatefulIngestionSourceBase):
     def get_workunits_internal(
         self,
     ) -> Iterable[MetadataWorkUnit]:
+        for metadata in self._get_metadata():
+            if isinstance(metadata, MetadataWorkUnit):
+                yield metadata
+            else:
+                yield from metadata.as_workunits()
+
+    def _get_metadata(self) -> Iterable[Union[MetadataWorkUnit, Entity]]:
         if not self.cassandra_api.authenticate():
             return
         keyspaces: List[CassandraKeyspace] = self.cassandra_api.get_keyspaces()
@@ -145,7 +145,7 @@ class CassandraSource(StatefulIngestionSourceBase):
                 self.report.report_dropped(keyspace_name)
                 continue
 
-            yield from self._generate_keyspace_container(keyspace)
+            yield self._generate_keyspace_container(keyspace)
 
             try:
                 yield from self._extract_tables_from_keyspace(keyspace_name)
@@ -170,21 +170,20 @@ class CassandraSource(StatefulIngestionSourceBase):
         if self.config.is_profiling_enabled():
             yield from self.profiler.get_workunits(self.cassandra_data)
 
-    def _generate_keyspace_container(
-        self, keyspace: CassandraKeyspace
-    ) -> Iterable[MetadataWorkUnit]:
+    def _generate_keyspace_container(self, keyspace: CassandraKeyspace) -> Container:
         keyspace_container_key = self._generate_keyspace_container_key(
             keyspace.keyspace_name
         )
-        yield from gen_containers(
-            container_key=keyspace_container_key,
-            name=keyspace.keyspace_name,
+
+        return Container(
+            keyspace_container_key,
+            display_name=keyspace.keyspace_name,
             qualified_name=keyspace.keyspace_name,
+            subtype=DatasetContainerSubTypes.KEYSPACE,
             extra_properties={
                 "durable_writes": str(keyspace.durable_writes),
                 "replication": json.dumps(keyspace.replication),
             },
-            sub_types=[DatasetContainerSubTypes.KEYSPACE],
         )
 
     def _generate_keyspace_container_key(self, keyspace_name: str) -> ContainerKey:
@@ -196,105 +195,55 @@ class CassandraSource(StatefulIngestionSourceBase):
         )
 
     # get all tables for a given keyspace, iterate over them to extract column metadata
-    def _extract_tables_from_keyspace(
-        self, keyspace_name: str
-    ) -> Iterable[MetadataWorkUnit]:
+    def _extract_tables_from_keyspace(self, keyspace_name: str) -> Iterable[Dataset]:
         self.cassandra_data.keyspaces.append(keyspace_name)
         tables: List[CassandraTable] = self.cassandra_api.get_tables(keyspace_name)
         for table in tables:
-            # define the dataset urn for this table to be used downstream
-            table_name: str = table.table_name
-            dataset_name: str = f"{keyspace_name}.{table_name}"
-
-            if not self.config.table_pattern.allowed(dataset_name):
-                self.report.report_dropped(dataset_name)
-                continue
-
-            self.cassandra_data.tables.setdefault(keyspace_name, []).append(table_name)
-            self.report.report_entity_scanned(dataset_name, ent_type="Table")
-
-            dataset_urn = make_dataset_urn_with_platform_instance(
-                platform=self.platform,
-                name=dataset_name,
-                env=self.config.env,
-                platform_instance=self.config.platform_instance,
+            dataset = self._generate_table(keyspace_name, table)
+            if dataset:
+                yield dataset
+
+    def _generate_table(
+        self, keyspace_name: str, table: CassandraTable
+    ) -> Optional[Dataset]:
+        table_name: str = table.table_name
+        dataset_name: str = f"{keyspace_name}.{table_name}"
+
+        self.report.report_entity_scanned(dataset_name, ent_type="Table")
+        if not self.config.table_pattern.allowed(dataset_name):
+            self.report.report_dropped(dataset_name)
+            return None
+
+        self.cassandra_data.tables.setdefault(keyspace_name, []).append(table_name)
+
+        schema_fields = None
+        try:
+            schema_fields = self._extract_columns_from_table(keyspace_name, table_name)
+        except Exception as e:
+            self.report.failure(
+                message="Failed to extract columns from table",
+                context=dataset_name,
+                exc=e,
             )
 
-            # 1. Extract columns from table, then construct and emit the schemaMetadata aspect.
-            try:
-                yield from self._extract_columns_from_table(
-                    keyspace_name, table_name, dataset_urn
-                )
-            except Exception as e:
-                self.report.failure(
-                    message="Failed to extract columns from table",
-                    context=table_name,
-                    exc=e,
-                )
-
-            yield MetadataChangeProposalWrapper(
-                entityUrn=dataset_urn,
-                aspect=StatusClass(removed=False),
-            ).as_workunit()
-
-            yield MetadataChangeProposalWrapper(
-                entityUrn=dataset_urn,
-                aspect=SubTypesClass(
-                    typeNames=[
-                        DatasetSubTypes.TABLE,
-                    ]
-                ),
-            ).as_workunit()
-
-            yield MetadataChangeProposalWrapper(
-                entityUrn=dataset_urn,
-                aspect=DatasetPropertiesClass(
-                    name=table_name,
-                    qualifiedName=f"{keyspace_name}.{table_name}",
-                    description=table.comment,
-                    customProperties={
-                        "bloom_filter_fp_chance": str(table.bloom_filter_fp_chance),
-                        "caching": json.dumps(table.caching),
-                        "compaction": json.dumps(table.compaction),
-                        "compression": json.dumps(table.compression),
-                        "crc_check_chance": str(table.crc_check_chance),
-                        "dclocal_read_repair_chance": str(
-                            table.dclocal_read_repair_chance
-                        ),
-                        "default_time_to_live": str(table.default_time_to_live),
-                        "extensions": json.dumps(table.extensions),
-                        "gc_grace_seconds": str(table.gc_grace_seconds),
-                        "max_index_interval": str(table.max_index_interval),
-                        "min_index_interval": str(table.min_index_interval),
-                        "memtable_flush_period_in_ms": str(
-                            table.memtable_flush_period_in_ms
-                        ),
-                        "read_repair_chance": str(table.read_repair_chance),
-                        "speculative_retry": str(table.speculative_retry),
-                    },
-                ),
-            ).as_workunit()
-
-            yield from add_dataset_to_container(
-                container_key=self._generate_keyspace_container_key(keyspace_name),
-                dataset_urn=dataset_urn,
-            )
-
-            if self.config.platform_instance:
-                yield MetadataChangeProposalWrapper(
-                    entityUrn=dataset_urn,
-                    aspect=DataPlatformInstanceClass(
-                        platform=make_data_platform_urn(self.platform),
-                        instance=make_dataplatform_instance_urn(
-                            self.platform, self.config.platform_instance
-                        ),
-                    ),
-                ).as_workunit()
+        return Dataset(
+            platform=self.platform,
+            name=dataset_name,
+            env=self.config.env,
+            platform_instance=self.config.platform_instance,
+            subtype=DatasetSubTypes.TABLE,
+            parent_container=self._generate_keyspace_container_key(keyspace_name),
+            schema=schema_fields,
+            display_name=table_name,
+            qualified_name=dataset_name,
+            description=table.comment,
+            custom_properties=self._get_dataset_custom_props(table),
+        )
 
     # get all columns for a given table, iterate over them to extract column metadata
     def _extract_columns_from_table(
-        self, keyspace_name: str, table_name: str, dataset_urn: str
-    ) -> Iterable[MetadataWorkUnit]:
+        self, keyspace_name: str, table_name: str
+    ) -> Optional[List[SchemaField]]:
         column_infos: List[CassandraColumn] = self.cassandra_api.get_columns(
             keyspace_name, table_name
         )
@@ -305,147 +254,117 @@ class CassandraSource(StatefulIngestionSourceBase):
             self.report.report_warning(
                 message="Table has no columns, skipping", context=table_name
             )
-            return
+            return None
 
+        # Tricky: we also save the column info to a global store.
         jsonable_column_infos: List[Dict[str, Any]] = []
         for column in column_infos:
             self.cassandra_data.columns.setdefault(table_name, []).append(column)
             jsonable_column_infos.append(dataclasses.asdict(column))
 
-        schema_metadata: SchemaMetadata = SchemaMetadata(
-            schemaName=table_name,
-            platform=make_data_platform_urn(self.platform),
-            version=0,
-            hash="",
-            platformSchema=OtherSchemaClass(
-                rawSchema=json.dumps(jsonable_column_infos)
-            ),
-            fields=schema_fields,
-        )
-
-        yield MetadataChangeProposalWrapper(
-            entityUrn=dataset_urn,
-            aspect=schema_metadata,
-        ).as_workunit()
+        return schema_fields
 
-    def _extract_views_from_keyspace(
-        self, keyspace_name: str
-    ) -> Iterable[MetadataWorkUnit]:
+    def _extract_views_from_keyspace(self, keyspace_name: str) -> Iterable[Dataset]:
         views: List[CassandraView] = self.cassandra_api.get_views(keyspace_name)
         for view in views:
-            view_name: str = view.view_name
-            dataset_name: str = f"{keyspace_name}.{view_name}"
-            self.report.report_entity_scanned(dataset_name)
-            dataset_urn: str = make_dataset_urn_with_platform_instance(
-                platform=self.platform,
-                name=dataset_name,
-                env=self.config.env,
-                platform_instance=self.config.platform_instance,
+            dataset = self._generate_view(keyspace_name, view)
+            if dataset:
+                yield dataset
+
+    def _generate_view(
+        self, keyspace_name: str, view: CassandraView
+    ) -> Optional[Dataset]:
+        view_name: str = view.view_name
+        dataset_name: str = f"{keyspace_name}.{view_name}"
+
+        self.report.report_entity_scanned(dataset_name, ent_type="View")
+        if not self.config.table_pattern.allowed(dataset_name):
+            # TODO: Maybe add a view_pattern instead of reusing table_pattern?
+            self.report.report_dropped(dataset_name)
+            return None
+
+        schema_fields = None
+        try:
+            schema_fields = self._extract_columns_from_table(keyspace_name, view_name)
+        except Exception as e:
+            self.report.failure(
+                message="Failed to extract columns from views",
+                context=view_name,
+                exc=e,
             )
 
-            yield MetadataChangeProposalWrapper(
-                entityUrn=dataset_urn,
-                aspect=StatusClass(removed=False),
-            ).as_workunit()
-
-            yield MetadataChangeProposalWrapper(
-                entityUrn=dataset_urn,
-                aspect=SubTypesClass(
-                    typeNames=[
-                        DatasetSubTypes.VIEW,
-                    ]
-                ),
-            ).as_workunit()
-
-            yield MetadataChangeProposalWrapper(
-                entityUrn=dataset_urn,
-                aspect=ViewPropertiesClass(
+        dataset = Dataset(
+            platform=self.platform,
+            name=dataset_name,
+            env=self.config.env,
+            platform_instance=self.config.platform_instance,
+            subtype=DatasetSubTypes.VIEW,
+            parent_container=self._generate_keyspace_container_key(keyspace_name),
+            schema=schema_fields,
+            display_name=view_name,
+            qualified_name=dataset_name,
+            description=view.comment,
+            custom_properties=self._get_dataset_custom_props(view),
+            extra_aspects=[
+                ViewPropertiesClass(
                     materialized=True,
                     viewLogic=view.where_clause,  # Use the WHERE clause as view logic
                     viewLanguage="CQL",  # Use "CQL" as the language
                 ),
-            ).as_workunit()
-
-            yield MetadataChangeProposalWrapper(
-                entityUrn=dataset_urn,
-                aspect=DatasetPropertiesClass(
-                    name=view_name,
-                    qualifiedName=f"{keyspace_name}.{view_name}",
-                    description=view.comment,
-                    customProperties={
-                        "bloom_filter_fp_chance": str(view.bloom_filter_fp_chance),
-                        "caching": json.dumps(view.caching),
-                        "compaction": json.dumps(view.compaction),
-                        "compression": json.dumps(view.compression),
-                        "crc_check_chance": str(view.crc_check_chance),
-                        "include_all_columns": str(view.include_all_columns),
-                        "dclocal_read_repair_chance": str(
-                            view.dclocal_read_repair_chance
-                        ),
-                        "default_time_to_live": str(view.default_time_to_live),
-                        "extensions": json.dumps(view.extensions),
-                        "gc_grace_seconds": str(view.gc_grace_seconds),
-                        "max_index_interval": str(view.max_index_interval),
-                        "min_index_interval": str(view.min_index_interval),
-                        "memtable_flush_period_in_ms": str(
-                            view.memtable_flush_period_in_ms
-                        ),
-                        "read_repair_chance": str(view.read_repair_chance),
-                        "speculative_retry": str(view.speculative_retry),
-                    },
-                ),
-            ).as_workunit()
+            ],
+        )
 
-            try:
-                yield from self._extract_columns_from_table(
-                    keyspace_name, view_name, dataset_urn
-                )
-            except Exception as e:
-                self.report.failure(
-                    message="Failed to extract columns from views",
-                    context=view_name,
-                    exc=e,
+        # Construct and emit lineage off of 'base_table_name'
+        # NOTE: we don't need to use 'base_table_id' since table is always in same keyspace, see https://docs.datastax.com/en/cql-oss/3.3/cql/cql_reference/cqlCreateMaterializedView.html#cqlCreateMaterializedView__keyspace-name
+        upstream_urn: str = make_dataset_urn_with_platform_instance(
+            platform=self.platform,
+            name=f"{keyspace_name}.{view.base_table_name}",
+            env=self.config.env,
+            platform_instance=self.config.platform_instance,
+        )
+        fineGrainedLineages = self.get_upstream_fields_of_field_in_datasource(
+            view_name, str(dataset.urn), upstream_urn
+        )
+        upstream_lineage = UpstreamLineageClass(
+            upstreams=[
+                UpstreamClass(
+                    dataset=upstream_urn,
+                    type=DatasetLineageTypeClass.VIEW,
                 )
+            ],
+            fineGrainedLineages=fineGrainedLineages,
+        )
 
-            # Construct and emit lineage off of 'base_table_name'
-            # NOTE: we don't need to use 'base_table_id' since table is always in same keyspace, see https://docs.datastax.com/en/cql-oss/3.3/cql/cql_reference/cqlCreateMaterializedView.html#cqlCreateMaterializedView__keyspace-name
-            upstream_urn: str = make_dataset_urn_with_platform_instance(
-                platform=self.platform,
-                name=f"{keyspace_name}.{view.table_name}",
-                env=self.config.env,
-                platform_instance=self.config.platform_instance,
-            )
-            fineGrainedLineages = self.get_upstream_fields_of_field_in_datasource(
-                view_name, dataset_urn, upstream_urn
-            )
-            yield MetadataChangeProposalWrapper(
-                entityUrn=dataset_urn,
-                aspect=UpstreamLineageClass(
-                    upstreams=[
-                        UpstreamClass(
-                            dataset=upstream_urn,
-                            type=DatasetLineageTypeClass.VIEW,
-                        )
-                    ],
-                    fineGrainedLineages=fineGrainedLineages,
-                ),
-            ).as_workunit()
-
-            yield from add_dataset_to_container(
-                container_key=self._generate_keyspace_container_key(keyspace_name),
-                dataset_urn=dataset_urn,
+        dataset.set_upstreams(upstream_lineage)
+
+        return dataset
+
+    def _get_dataset_custom_props(
+        self, dataset: CassandraSharedDatasetFields
+    ) -> Dict[str, str]:
+        props = {
+            "bloom_filter_fp_chance": str(dataset.bloom_filter_fp_chance),
+            "caching": json.dumps(dataset.caching),
+            "compaction": json.dumps(dataset.compaction),
+            "compression": json.dumps(dataset.compression),
+            "crc_check_chance": str(dataset.crc_check_chance),
+            "dclocal_read_repair_chance": str(dataset.dclocal_read_repair_chance),
+            "default_time_to_live": str(dataset.default_time_to_live),
+            "extensions": json.dumps(dataset.extensions),
+            "gc_grace_seconds": str(dataset.gc_grace_seconds),
+            "max_index_interval": str(dataset.max_index_interval),
+            "min_index_interval": str(dataset.min_index_interval),
+            "memtable_flush_period_in_ms": str(dataset.memtable_flush_period_in_ms),
+            "read_repair_chance": str(dataset.read_repair_chance),
+            "speculative_retry": str(dataset.speculative_retry),
+        }
+        if isinstance(dataset, CassandraView):
+            props.update(
+                {
+                    "include_all_columns": str(dataset.include_all_columns),
+                }
             )
-
-            if self.config.platform_instance:
-                yield MetadataChangeProposalWrapper(
-                    entityUrn=dataset_urn,
-                    aspect=DataPlatformInstanceClass(
-                        platform=make_data_platform_urn(self.platform),
-                        instance=make_dataplatform_instance_urn(
-                            self.platform, self.config.platform_instance
-                        ),
-                    ),
-                ).as_workunit()
+        return props
 
     def get_upstream_fields_of_field_in_datasource(
         self, table_name: str, dataset_urn: str, upstream_urn: str

datahub/ingestion/source/cassandra/cassandra_api.py

@@ -23,9 +23,9 @@ class CassandraKeyspace:
 
 
 @dataclass
-class CassandraTable:
+class CassandraSharedDatasetFields:
     keyspace_name: str
-    table_name: str
+
     bloom_filter_fp_chance: Optional[float]
     caching: Optional[Dict[str, str]]
     comment: Optional[str]
@@ -43,6 +43,11 @@ class CassandraTable:
     speculative_retry: Optional[str]
 
 
+@dataclass
+class CassandraTable(CassandraSharedDatasetFields):
+    table_name: str
+
+
 @dataclass
 class CassandraColumn:
     keyspace_name: str
@@ -55,8 +60,10 @@ class CassandraColumn:
 
 
 @dataclass
-class CassandraView(CassandraTable):
+class CassandraView(CassandraSharedDatasetFields):
     view_name: str
+
+    base_table_name: str
     include_all_columns: Optional[bool]
     where_clause: str = ""
 
@@ -261,7 +268,7 @@ class CassandraAPI:
             views = self.get(CassandraQueries.GET_VIEWS_QUERY, [keyspace_name])
             view_list = [
                 CassandraView(
-                    table_name=row.base_table_name,
+                    base_table_name=row.base_table_name,
                     keyspace_name=row.keyspace_name,
                     view_name=row.view_name,
                     bloom_filter_fp_chance=row.bloom_filter_fp_chance,

datahub/ingestion/source/delta_lake/config.py

@@ -13,6 +13,9 @@ from datahub.configuration.source_common import (
 )
 from datahub.ingestion.source.aws.aws_common import AwsConnectionConfig
 from datahub.ingestion.source.aws.s3_util import is_s3_uri
+from datahub.ingestion.source.state.stateful_ingestion_base import (
+    StatefulIngestionConfigBase,
+)
 
 # hide annoying debug errors from py4j
 logging.getLogger("py4j").setLevel(logging.ERROR)
@@ -35,7 +38,11 @@ class S3(ConfigModel):
     )
 
 
-class DeltaLakeSourceConfig(PlatformInstanceConfigMixin, EnvConfigMixin):
+class DeltaLakeSourceConfig(
+    PlatformInstanceConfigMixin,
+    EnvConfigMixin,
+    StatefulIngestionConfigBase,
+):
     base_path: str = Field(
         description="Path to table (s3 or local file system). If path is not a delta table path "
         "then all subfolders will be scanned to detect and ingest delta tables."

datahub/ingestion/source/delta_lake/report.py

@@ -1,12 +1,14 @@
 import dataclasses
 from dataclasses import field as dataclass_field
 
-from datahub.ingestion.api.source import SourceReport
+from datahub.ingestion.source.state.stale_entity_removal_handler import (
+    StaleEntityRemovalSourceReport,
+)
 from datahub.utilities.lossy_collections import LossyList
 
 
 @dataclasses.dataclass
-class DeltaLakeSourceReport(SourceReport):
+class DeltaLakeSourceReport(StaleEntityRemovalSourceReport):
     files_scanned = 0
     filtered: LossyList[str] = dataclass_field(default_factory=LossyList)
 

datahub/ingestion/source/delta_lake/source.py

@@ -2,7 +2,7 @@ import json
 import logging
 import os
 import time
-from typing import Dict, Iterable, List
+from typing import Dict, Iterable, List, Optional
 from urllib.parse import urlparse
 
 from deltalake import DeltaTable
@@ -21,7 +21,7 @@ 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.aws.s3_boto_utils import get_s3_tags
 from datahub.ingestion.source.aws.s3_util import (
@@ -36,6 +36,12 @@ from datahub.ingestion.source.delta_lake.delta_lake_utils import (
     read_delta_table,
 )
 from datahub.ingestion.source.delta_lake.report import DeltaLakeSourceReport
+from datahub.ingestion.source.state.stale_entity_removal_handler import (
+    StaleEntityRemovalHandler,
+)
+from datahub.ingestion.source.state.stateful_ingestion_base import (
+    StatefulIngestionSourceBase,
+)
 from datahub.metadata.com.linkedin.pegasus2avro.common import Status
 from datahub.metadata.com.linkedin.pegasus2avro.metadata.snapshot import DatasetSnapshot
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
@@ -79,7 +85,7 @@ OPERATION_STATEMENT_TYPES = {
 @config_class(DeltaLakeSourceConfig)
 @support_status(SupportStatus.INCUBATING)
 @capability(SourceCapability.TAGS, "Can extract S3 object/bucket tags if enabled")
-class DeltaLakeSource(Source):
+class DeltaLakeSource(StatefulIngestionSourceBase):
     """
     This plugin extracts:
     - Column types and schema associated with each delta table
@@ -100,9 +106,10 @@ class DeltaLakeSource(Source):
     storage_options: Dict[str, str]
 
     def __init__(self, config: DeltaLakeSourceConfig, ctx: PipelineContext):
-        super().__init__(ctx)
+        super().__init__(config, ctx)
+        self.ctx = ctx
         self.source_config = config
-        self.report = DeltaLakeSourceReport()
+        self.report: DeltaLakeSourceReport = DeltaLakeSourceReport()
         if self.source_config.is_s3:
             if (
                 self.source_config.s3 is None
@@ -331,6 +338,14 @@ class DeltaLakeSource(Source):
         for folder in os.listdir(path):
             yield os.path.join(path, folder)
 
+    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]:
         self.container_WU_creator = ContainerWUCreator(
             self.source_config.platform,