acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/mongodb.py

@@ -36,7 +36,10 @@ from datahub.ingestion.api.decorators import (
 )
 from datahub.ingestion.api.source import MetadataWorkUnitProcessor
 from datahub.ingestion.api.workunit import MetadataWorkUnit
-from datahub.ingestion.source.common.subtypes import DatasetContainerSubTypes
+from datahub.ingestion.source.common.subtypes import (
+    DatasetContainerSubTypes,
+    SourceCapabilityModifier,
+)
 from datahub.ingestion.source.schema_inference.object import (
     SchemaDescription,
     construct_schema,
@@ -249,6 +252,13 @@ def construct_schema_pymongo(
 @support_status(SupportStatus.CERTIFIED)
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(SourceCapability.SCHEMA_METADATA, "Enabled by default")
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.DATABASE,
+    ],
+)
 @dataclass
 class MongoDBSource(StatefulIngestionSourceBase):
     """

datahub/ingestion/source/neo4j/neo4j_source.py

@@ -1,74 +1,56 @@
 import logging
-import time
 from dataclasses import dataclass
-from typing import Any, Dict, Iterable, List, Optional, Type, Union
+from typing import Dict, Iterable, List, Optional, Tuple
 
 import pandas as pd
 from neo4j import GraphDatabase
-from pydantic.fields import Field
+from pydantic import Field
 
 from datahub.configuration.source_common import (
     EnvConfigMixin,
+    PlatformInstanceConfigMixin,
 )
-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
 from datahub.ingestion.api.decorators import (
     SupportStatus,
+    capability,
     config_class,
     platform_name,
     support_status,
 )
 from datahub.ingestion.api.source import (
     MetadataWorkUnitProcessor,
+    SourceCapability,
 )
 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,
+    StatefulStaleMetadataRemovalConfig,
 )
 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,
-    BooleanTypeClass,
-    DatasetPropertiesClass,
-    DateTypeClass,
-    NullTypeClass,
-    NumberTypeClass,
-    OtherSchemaClass,
-    SchemaFieldClass,
-    SchemaMetadataClass,
-    StringTypeClass,
-    SubTypesClass,
-    UnionTypeClass,
-)
+from datahub.sdk.dataset import Dataset
 
 log = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
 
-_type_mapping: Dict[Union[Type, str], Type] = {
-    "list": UnionTypeClass,
-    "boolean": BooleanTypeClass,
-    "integer": NumberTypeClass,
-    "local_date_time": DateTypeClass,
-    "float": NumberTypeClass,
-    "string": StringTypeClass,
-    "date": DateTypeClass,
-    "node": StringTypeClass,
-    "relationship": StringTypeClass,
-}
+# Neo4j object types
+_NODE = "node"
+_RELATIONSHIP = "relationship"
 
 
-class Neo4jConfig(EnvConfigMixin, StatefulIngestionConfigBase):
+class Neo4jConfig(
+    StatefulIngestionConfigBase, EnvConfigMixin, PlatformInstanceConfigMixin
+):
     username: str = Field(description="Neo4j Username")
     password: str = Field(description="Neo4j Password")
     uri: str = Field(description="The URI for the Neo4j server")
-    env: str = Field(description="Neo4j env")
+
+    stateful_ingestion: Optional[StatefulStaleMetadataRemovalConfig] = None
 
 
 @dataclass
@@ -79,114 +61,111 @@ class Neo4jSourceReport(StatefulIngestionReport):
 
 @platform_name("Neo4j", id="neo4j")
 @config_class(Neo4jConfig)
+@capability(
+    SourceCapability.PLATFORM_INSTANCE, "Supported via the `platform_instance` config"
+)
 @support_status(SupportStatus.CERTIFIED)
 class Neo4jSource(StatefulIngestionSourceBase):
-    NODE = "node"
-    RELATIONSHIP = "relationship"
-    PLATFORM = "neo4j"
+    config: Neo4jConfig
+    report: Neo4jSourceReport
 
-    def __init__(self, ctx: PipelineContext, config: Neo4jConfig):
+    def __init__(self, config: Neo4jConfig, ctx: PipelineContext):
+        super().__init__(config, ctx)
         self.ctx = ctx
         self.config = config
+        self.platform = "neo4j"
         self.report: Neo4jSourceReport = Neo4jSourceReport()
 
     @classmethod
-    def create(cls, config_dict, ctx):
+    def create(cls, config_dict: Dict, ctx: PipelineContext) -> "Neo4jSource":
         config = Neo4jConfig.parse_obj(config_dict)
-        return cls(ctx, config)
-
-    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())
-
-    def get_schema_field_class(
-        self, col_name: str, col_type: str, **kwargs: Any
-    ) -> SchemaFieldClass:
-        if kwargs["obj_type"] == self.NODE and col_type == self.RELATIONSHIP:
-            col_type = self.NODE
-        else:
-            col_type = col_type
-        return SchemaFieldClass(
-            fieldPath=col_name,
-            type=self.get_field_type(col_type),
-            nativeDataType=col_type,
-            description=col_type.upper()
-            if col_type in (self.NODE, self.RELATIONSHIP)
-            else col_type,
-            lastModified=AuditStampClass(
-                time=round(time.time() * 1000), actor="urn:li:corpuser:ingestion"
-            ),
+        return cls(config, ctx)
+
+    def create_schema_field_tuple(
+        self, col_name: str, col_type: str, obj_type: Optional[str]
+    ) -> Tuple[str, str, str]:
+        """Convert Neo4j property to (field_name, field_type, description) tuple."""
+        # Special case: when a node has a relationship-typed property, treat it as a node reference
+        # This ensures relationship properties within nodes are described as "NODE" rather than "RELATIONSHIP"
+        column_type = (
+            _NODE if obj_type == _NODE and col_type == _RELATIONSHIP else col_type
         )
 
-    def add_properties(
+        description = (
+            column_type.upper()
+            if column_type in (_NODE, _RELATIONSHIP)
+            else column_type
+        )
+
+        return (col_name, column_type, description)
+
+    def get_subtype_from_obj_type(self, obj_type: str) -> str:
+        """Map Neo4j object type to DataHub subtype."""
+        if obj_type == _NODE:
+            return DatasetSubTypes.NEO4J_NODE
+        elif obj_type == _RELATIONSHIP:
+            return DatasetSubTypes.NEO4J_RELATIONSHIP
+        return DatasetSubTypes.NEO4J_NODE  # default fallback
+
+    def create_neo4j_dataset(
         self,
         dataset: str,
+        columns: list,
+        obj_type: Optional[str] = None,
         description: Optional[str] = None,
-        custom_properties: Optional[Dict[str, str]] = None,
-    ) -> MetadataChangeProposalWrapper:
-        dataset_properties = DatasetPropertiesClass(
-            description=description,
-            customProperties=custom_properties,
-        )
-        return MetadataChangeProposalWrapper(
-            entityUrn=make_dataset_urn(
-                platform=self.PLATFORM, name=dataset, env=self.config.env
-            ),
-            aspect=dataset_properties,
-        )
-
-    def generate_neo4j_object(
-        self, dataset: str, columns: list, obj_type: Optional[str] = None
-    ) -> MetadataChangeProposalWrapper:
+    ) -> Optional[Dataset]:
+        """Create Dataset entity with Neo4j schema and metadata."""
         try:
-            fields = [
-                self.get_schema_field_class(key, value.lower(), obj_type=obj_type)
+            schema_fields = [
+                self.create_schema_field_tuple(
+                    col_name=key, col_type=value.lower(), obj_type=obj_type
+                )
                 for d in columns
                 for key, value in d.items()
             ]
-            mcp = MetadataChangeProposalWrapper(
-                entityUrn=make_dataset_urn(
-                    platform=self.PLATFORM, name=dataset, env=self.config.env
-                ),
-                aspect=SchemaMetadataClass(
-                    schemaName=dataset,
-                    platform=make_data_platform_urn(self.PLATFORM),
-                    version=0,
-                    hash="",
-                    platformSchema=OtherSchemaClass(rawSchema=""),
-                    lastModified=AuditStampClass(
-                        time=round(time.time() * 1000),
-                        actor="urn:li:corpuser:ingestion",
-                    ),
-                    fields=fields,
-                ),
+
+            return Dataset(
+                platform=self.platform,
+                name=dataset,
+                platform_instance=self.config.platform_instance,
+                env=self.config.env,
+                schema=schema_fields,
+                subtype=self.get_subtype_from_obj_type(obj_type or _NODE),
+                description=description,
             )
-            self.report.obj_created += 1
+
         except Exception as e:
             log.error(e)
-            self.report.obj_failures += 1
-        return mcp
+            self.report.report_failure(
+                message="Failed to process dataset",
+                context=dataset,
+                exc=e,
+            )
+            return None
 
-    def get_neo4j_metadata(self, query: str) -> pd.DataFrame:
+    def get_neo4j_metadata(self, query: str) -> Optional[pd.DataFrame]:
         driver = GraphDatabase.driver(
             self.config.uri, auth=(self.config.username, self.config.password)
         )
         """
-        This process retrieves the metadata for Neo4j objects using an APOC query, which returns a dictionary
-        with two columns: key and value. The key represents the Neo4j object, while the value contains the
-        corresponding metadata.
+        This process retrieves the metadata for Neo4j objects using an APOC query,
+        which returns a dictionary with two columns: key and value. The key represents
+        the Neo4j object, while the value contains the corresponding metadata.
 
-        When data is returned from Neo4j, much of the relationship metadata is stored with the relevant node's
-        metadata. Consequently, the objects are organized into two separate dataframes: one for nodes and one for
-        relationships.
+        When data is returned from Neo4j, much of the relationship metadata is stored
+        with the relevant node's metadata. Consequently, the objects are organized
+        into two separate dataframes: one for nodes and one for relationships.
 
-        In the node dataframe, several fields are extracted and added as new columns. Similarly, in the relationship
-        dataframe, certain fields are parsed out, while others require metadata from the nodes dataframe.
+        In the node dataframe, several fields are extracted and added as new columns.
+        Similarly, in the relationship dataframe, certain fields are parsed out,
+        while others require metadata from the nodes dataframe.
 
-        Once the data is parsed and these two dataframes are created, we combine a subset of their columns into a
-        single dataframe, which will be used to create the DataHub objects.
+        Once the data is parsed and these two dataframes are created, we combine
+        a subset of their columns into a single dataframe, which will be used to
+        create the DataHub objects.
 
-        See the docs for examples of metadata:  metadata-ingestion/docs/sources/neo4j/neo4j.md
+        See the docs for examples of metadata:
+        metadata-ingestion/docs/sources/neo4j/neo4j.md
         """
         try:
             log.info(f"{query}")
@@ -201,16 +180,17 @@ class Neo4jSource(StatefulIngestionSourceBase):
 
             union_cols = ["key", "obj_type", "property_data_types", "description"]
             df = pd.concat([node_df[union_cols], rel_df[union_cols]])
+            return df
         except Exception as e:
             self.report.failure(
                 message="Failed to get neo4j metadata",
                 exc=e,
             )
 
-        return df
+        return None
 
     def process_nodes(self, data: list) -> pd.DataFrame:
-        nodes = [record for record in data if record["value"]["type"] == self.NODE]
+        nodes = [record for record in data if record["value"]["type"] == _NODE]
         node_df = pd.DataFrame(
             nodes,
             columns=["key", "value"],
@@ -233,9 +213,7 @@ class Neo4jSource(StatefulIngestionSourceBase):
         return node_df
 
     def process_relationships(self, data: list, node_df: pd.DataFrame) -> pd.DataFrame:
-        rels = [
-            record for record in data if record["value"]["type"] == self.RELATIONSHIP
-        ]
+        rels = [record for record in data if record["value"]["type"] == _RELATIONSHIP]
         rel_df = pd.DataFrame(rels, columns=["key", "value"])
         rel_df["obj_type"] = rel_df["value"].apply(
             lambda record: self.get_obj_type(record)
@@ -303,49 +281,40 @@ class Neo4jSource(StatefulIngestionSourceBase):
         ]
 
     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;"
+        query = (
+            "CALL apoc.meta.schema() YIELD value UNWIND keys(value) AS key "
+            "RETURN key, value[key] AS value;"
         )
+        df = self.get_neo4j_metadata(query)
+        if df is None:
+            log.warning("No metadata retrieved from Neo4j")
+            return
+
         for _, row in df.iterrows():
             try:
-                yield MetadataWorkUnit(
-                    id=row["key"],
-                    mcp=self.generate_neo4j_object(
-                        columns=row["property_data_types"],
-                        dataset=row["key"],
-                    ),
-                    is_primary_source=True,
+                dataset_obj = self.create_neo4j_dataset(
+                    dataset=row["key"],
+                    columns=row["property_data_types"],
+                    obj_type=row["obj_type"],
+                    description=row["description"],
                 )
 
-                yield MetadataWorkUnit(
-                    id=row["key"],
-                    mcp=MetadataChangeProposalWrapper(
-                        entityUrn=make_dataset_urn(
-                            platform=self.PLATFORM,
-                            name=row["key"],
-                            env=self.config.env,
-                        ),
-                        aspect=SubTypesClass(
-                            typeNames=[
-                                DatasetSubTypes.NEO4J_NODE
-                                if row["obj_type"] == self.NODE
-                                else DatasetSubTypes.NEO4J_RELATIONSHIP
-                            ]
-                        ),
-                    ),
-                )
-
-                yield MetadataWorkUnit(
-                    id=row["key"],
-                    mcp=self.add_properties(
-                        dataset=row["key"],
-                        custom_properties=None,
-                        description=row["description"],
-                    ),
-                )
+                if dataset_obj:
+                    yield from dataset_obj.as_workunits()
+                    self.report.obj_created += 1
+                else:
+                    log.warning(f"Failed to create dataset object for {row['key']}")
+                    self.report.obj_failures += 1
 
             except Exception as e:
-                raise e
+                log.warning(f"Failed to process row {row['key']}: {str(e)}")
+                self.report.report_warning(
+                    title="Error processing Neo4j metadata",
+                    message="Some entities will be missed",
+                    context=row["key"],
+                    exc=e,
+                )
+                self.report.obj_failures += 1
 
-    def get_report(self):
+    def get_report(self) -> "Neo4jSourceReport":
         return self.report

datahub/ingestion/source/nifi.py

@@ -72,7 +72,7 @@ NIFI = "nifi"
 # and here - https://github.com/psf/requests/issues/1573
 class SSLAdapter(HTTPAdapter):
     def __init__(self, certfile, keyfile, password=None):
-        self.context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
+        self.context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
         self.context.load_cert_chain(
             certfile=certfile, keyfile=keyfile, password=password
         )
@@ -166,7 +166,7 @@ class NifiSourceConfig(StatefulIngestionConfigBase, EnvConfigMixin):
     )
 
     @root_validator(skip_on_failure=True)
-    def validate_auth_params(cla, values):
+    def validate_auth_params(cls, values):
         if values.get("auth") is NifiAuthType.CLIENT_CERT and not values.get(
             "client_cert_file"
         ):
@@ -703,7 +703,7 @@ class NifiSource(StatefulIngestionSourceBase):
             if (
                 component.nifi_type is NifiType.PROCESSOR
                 and component.type
-                not in NifiProcessorProvenanceEventAnalyzer.KNOWN_INGRESS_EGRESS_PROCESORS.keys()
+                not in NifiProcessorProvenanceEventAnalyzer.KNOWN_INGRESS_EGRESS_PROCESORS
             ) or component.nifi_type not in [
                 NifiType.PROCESSOR,
                 NifiType.REMOTE_INPUT_PORT,
@@ -977,7 +977,7 @@ class NifiSource(StatefulIngestionSourceBase):
                     )
 
             for incoming_from in incoming:
-                if incoming_from in self.nifi_flow.remotely_accessible_ports.keys():
+                if incoming_from in self.nifi_flow.remotely_accessible_ports:
                     dataset_name = f"{self.config.site_name}.{self.nifi_flow.remotely_accessible_ports[incoming_from].name}"
                     dataset_urn = builder.make_dataset_urn(
                         NIFI, dataset_name, self.config.env
@@ -994,7 +994,7 @@ class NifiSource(StatefulIngestionSourceBase):
                     )
 
             for outgoing_to in outgoing:
-                if outgoing_to in self.nifi_flow.remotely_accessible_ports.keys():
+                if outgoing_to in self.nifi_flow.remotely_accessible_ports:
                     dataset_name = f"{self.config.site_name}.{self.nifi_flow.remotely_accessible_ports[outgoing_to].name}"
                     dataset_urn = builder.make_dataset_urn(
                         NIFI, dataset_name, self.config.env