update based on feedback and upstream changes

Author: acrylJonny

metadata-ingestion/src/datahub/ingestion/source/neo4j/neo4j_source.py

@@ -32,7 +32,6 @@
 from datahub.ingestion.source.common.subtypes import DatasetSubTypes
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
-    StaleEntityRemovalSourceReport,
     StatefulStaleMetadataRemovalConfig,
 )
 from datahub.ingestion.source.state.stateful_ingestion_base import (
@@ -75,7 +74,6 @@
 class Neo4jConfig(
     StatefulIngestionConfigBase, EnvConfigMixin, PlatformInstanceConfigMixin
 ):
-    platform: str = Field(default="neo4j", hidden_from_docs=True)
     username: str = Field(description="Neo4j Username")
     password: str = Field(description="Neo4j Password")
     uri: str = Field(description="The URI for the Neo4j server")
@@ -92,7 +90,9 @@ class Neo4jSourceReport(StatefulIngestionReport):
 
 @platform_name("Neo4j", id="neo4j")
 @config_class(Neo4jConfig)
-@capability(SourceCapability.PLATFORM_INSTANCE, "Supported via the `platform_instance` config")
+@capability(
+    SourceCapability.PLATFORM_INSTANCE, "Supported via the `platform_instance` config"
+)
 @support_status(SupportStatus.CERTIFIED)
 class Neo4jSource(StatefulIngestionSourceBase):
     NODE = "node"
@@ -104,24 +104,14 @@ def __init__(self, config: Neo4jConfig, ctx: PipelineContext):
         super().__init__(config, ctx)
         self.ctx = ctx
         self.config = config
-        self.platform = self.config.platform
-        self.platform_instance = self.config.platform_instance
-        self.env = self.config.env
+        self.platform = "neo4j"
         self.report: Neo4jSourceReport = Neo4jSourceReport()
 
     @classmethod
     def create(cls, config_dict: Dict, ctx: PipelineContext) -> "Neo4jSource":
         config = Neo4jConfig.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.config, self.ctx
-            ).workunit_processor,
-        ]
-
     def get_field_type(self, attribute_type: Union[type, str]) -> SchemaFieldDataType:
         type_class: type = _type_mapping.get(attribute_type, NullTypeClass)
         return SchemaFieldDataType(type=type_class())
@@ -159,8 +149,8 @@ def add_properties(
             entityUrn=make_dataset_urn_with_platform_instance(
                 platform=self.platform,
                 name=dataset,
-                platform_instance=self.platform_instance,
-                env=self.env,
+                platform_instance=self.config.platform_instance,
+                env=self.config.env,
             ),
             aspect=dataset_properties,
         ).as_workunit()
@@ -178,8 +168,8 @@ def generate_neo4j_object(
                 entityUrn=make_dataset_urn_with_platform_instance(
                     platform=self.platform,
                     name=dataset,
-                    platform_instance=self.platform_instance,
-                    env=self.env,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
                 ),
                 aspect=SchemaMetadataClass(
                     schemaName=dataset,
@@ -359,8 +349,8 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
                     entityUrn=make_dataset_urn_with_platform_instance(
                         platform=self.platform,
                         name=row["key"],
-                        platform_instance=self.platform_instance,
-                        env=self.env,
+                        platform_instance=self.config.platform_instance,
+                        env=self.config.env,
                     ),
                     aspect=SubTypesClass(
                         typeNames=[

metadata-ingestion/tests/unit/test_neo4j_source.py

@@ -328,7 +328,7 @@ def test_get_relationships(source, mock_data):
     assert list(source.get_relationships(results[0]["value"]).keys()) == [
         "RELATIONSHIP_1"
     ]
-    assert source.get_relationships(results[2]["value"]) is None
+    assert source.get_relationships(results[2]["value"]) == {}
 
 
 def test_get_field_type(source):
@@ -384,7 +384,6 @@ def test_type_mapping(source, test_input, expected):
 
 def test_platform_instance_config(source_with_platform_instance):
     """Test platform instance configuration"""
-    assert source_with_platform_instance.platform_instance == "test-instance"
     assert source_with_platform_instance.config.platform_instance == "test-instance"
 
 
@@ -454,7 +453,6 @@ def test_default_values():
     config = Neo4jConfig(
         uri="neo4j://localhost:7687", username="neo4j", password="password", env="PROD"
     )
-    assert config.platform == "neo4j"
     assert config.platform_instance is None
     assert config.stateful_ingestion is None