acryl-datahub 1.2.0.8rc4__py3-none-any.whl → 1.2.0.9__py3-none-any.whl

datahub/ingestion/source/neo4j/neo4j_source.py

@@ -1,7 +1,6 @@
 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
@@ -11,11 +10,6 @@ from datahub.configuration.source_common import (
     EnvConfigMixin,
     PlatformInstanceConfigMixin,
 )
-from datahub.emitter.mce_builder import (
-    make_data_platform_urn,
-    make_dataset_urn_with_platform_instance,
-)
-from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SupportStatus,
@@ -28,7 +22,6 @@ from datahub.ingestion.api.source import (
     MetadataWorkUnitProcessor,
     SourceCapability,
 )
-from datahub.ingestion.api.source_helpers import auto_workunit
 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 (
@@ -40,36 +33,14 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
     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(
@@ -78,7 +49,6 @@ class Neo4jConfig(
     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
 
@@ -96,8 +66,6 @@ class Neo4jSourceReport(StatefulIngestionReport):
 )
 @support_status(SupportStatus.CERTIFIED)
 class Neo4jSource(StatefulIngestionSourceBase):
-    NODE = "node"
-    RELATIONSHIP = "relationship"
     config: Neo4jConfig
     report: Neo4jSourceReport
 
@@ -113,78 +81,59 @@ class Neo4jSource(StatefulIngestionSourceBase):
         config = Neo4jConfig.parse_obj(config_dict)
         return cls(config, ctx)
 
-    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"
-            ),
+    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
+        )
+
+        description = (
+            column_type.upper()
+            if column_type in (_NODE, _RELATIONSHIP)
+            else column_type
         )
 
-    def add_properties(
+        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,
-    ) -> Iterable[MetadataWorkUnit]:
-        dataset_properties = DatasetPropertiesClass(
-            description=description,
-            customProperties=custom_properties,
-        )
-        yield MetadataChangeProposalWrapper(
-            entityUrn=make_dataset_urn_with_platform_instance(
-                platform=self.platform,
-                name=dataset,
-                platform_instance=self.config.platform_instance,
-                env=self.config.env,
-            ),
-            aspect=dataset_properties,
-        ).as_workunit()
-
-    def generate_neo4j_object(
-        self, dataset: str, columns: list, obj_type: Optional[str] = None
-    ) -> Optional[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()
             ]
-            return MetadataChangeProposalWrapper(
-                entityUrn=make_dataset_urn_with_platform_instance(
-                    platform=self.platform,
-                    name=dataset,
-                    platform_instance=self.config.platform_instance,
-                    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,
             )
+
         except Exception as e:
             log.error(e)
             self.report.report_failure(
@@ -199,21 +148,24 @@ class Neo4jSource(StatefulIngestionSourceBase):
             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}")
@@ -238,7 +190,7 @@ class Neo4jSource(StatefulIngestionSourceBase):
         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"],
@@ -261,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)
@@ -331,51 +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:
-                neo4j_obj = self.generate_neo4j_object(
-                    columns=row["property_data_types"],
-                    dataset=row["key"],
-                )
-                if neo4j_obj:
-                    yield from auto_workunit([neo4j_obj])
-
-                yield MetadataChangeProposalWrapper(
-                    entityUrn=make_dataset_urn_with_platform_instance(
-                        platform=self.platform,
-                        name=row["key"],
-                        platform_instance=self.config.platform_instance,
-                        env=self.config.env,
-                    ),
-                    aspect=SubTypesClass(
-                        typeNames=[
-                            DatasetSubTypes.NEO4J_NODE
-                            if row["obj_type"] == self.NODE
-                            else DatasetSubTypes.NEO4J_RELATIONSHIP
-                        ]
-                    ),
-                ).as_workunit()
-
-                yield from self.add_properties(
+                dataset_obj = self.create_neo4j_dataset(
                     dataset=row["key"],
-                    custom_properties=None,
+                    columns=row["property_data_types"],
+                    obj_type=row["obj_type"],
                     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:
-                log.error(f"Failed to process row {row['key']}: {str(e)}")
-                self.report.report_failure(
-                    message="Error processing Neo4j metadata",
+                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) -> "Neo4jSourceReport":
         return self.report

datahub/ingestion/source/sql/sql_types.py

@@ -459,6 +459,25 @@ VERTICA_SQL_TYPES_MAP: Dict[str, Any] = {
     "uuid": StringType,
 }
 
+# Neo4j property types mapping
+# https://neo4j.com/docs/cypher-manual/current/values-and-types/property-structural-constructed/
+NEO4J_TYPES_MAP: Dict[str, Any] = {
+    "boolean": BooleanType,
+    "date": DateType,
+    "duration": TimeType,  # Neo4j duration represents a temporal amount
+    "float": NumberType,
+    "integer": NumberType,
+    "list": ArrayType,
+    "local_date_time": TimeType,
+    "local_time": TimeType,
+    "point": StringType,  # Neo4j point - spatial coordinate, represented as string
+    "string": StringType,
+    "zoned_date_time": TimeType,
+    "zoned_time": TimeType,
+    "node": StringType,  # Neo4j object type
+    "relationship": StringType,  # Neo4j object type
+}
+
 
 _merged_mapping = {
     "boolean": BooleanType,
@@ -478,6 +497,7 @@ _merged_mapping = {
     **TRINO_SQL_TYPES_MAP,
     **ATHENA_SQL_TYPES_MAP,
     **VERTICA_SQL_TYPES_MAP,
+    **NEO4J_TYPES_MAP,
 }
 
 
@@ -487,6 +507,8 @@ def resolve_sql_type(
 ) -> Optional[DATAHUB_FIELD_TYPE]:
     # In theory, we should use the platform-specific mapping where available.
     # However, the types don't ever conflict, so the merged mapping is fine.
+    # Wrong assumption - there ARE conflicts as the test_type_conflicts_across_platforms in test_sql_types.py shows.
+    # TODO: revisit this and make platform-specific mappings work.
     TypeClass: Optional[Type[DATAHUB_FIELD_TYPE]] = (
         _merged_mapping.get(column_type) if column_type else None
     )

datahub/metadata/_internal_schema_classes.py

@@ -23033,7 +23033,7 @@ class VersionBasedRetentionClass(DictWrapper):
     
     
 class ActorsClass(_Aspect):
-    """Provisioned users of a role"""
+    """Provisioned users and groups of a role"""
 
 
     ASPECT_NAME = 'actors'
@@ -23042,13 +23042,16 @@ class ActorsClass(_Aspect):
 
     def __init__(self,
         users: Union[None, List["RoleUserClass"]]=None,
+        groups: Union[None, List["RoleGroupClass"]]=None,
     ):
         super().__init__()
         
         self.users = users
+        self.groups = groups
     
     def _restore_defaults(self) -> None:
         self.users = self.RECORD_SCHEMA.fields_dict["users"].default
+        self.groups = self.RECORD_SCHEMA.fields_dict["groups"].default
     
     
     @property
@@ -23061,6 +23064,41 @@ class ActorsClass(_Aspect):
         self._inner_dict['users'] = value
     
     
+    @property
+    def groups(self) -> Union[None, List["RoleGroupClass"]]:
+        """List of provisioned groups of a role"""
+        return self._inner_dict.get('groups')  # type: ignore
+    
+    @groups.setter
+    def groups(self, value: Union[None, List["RoleGroupClass"]]) -> None:
+        self._inner_dict['groups'] = value
+    
+    
+class RoleGroupClass(DictWrapper):
+    """Provisioned groups of a role"""
+    
+    RECORD_SCHEMA = get_schema_type("com.linkedin.pegasus2avro.role.RoleGroup")
+    def __init__(self,
+        group: str,
+    ):
+        super().__init__()
+        
+        self.group = group
+    
+    def _restore_defaults(self) -> None:
+        self.group = str()
+    
+    
+    @property
+    def group(self) -> str:
+        """Link provisioned corp group for a role"""
+        return self._inner_dict.get('group')  # type: ignore
+    
+    @group.setter
+    def group(self, value: str) -> None:
+        self._inner_dict['group'] = value
+    
+    
 class RolePropertiesClass(_Aspect):
     """Information about a ExternalRoleProperties"""
 
@@ -27551,6 +27589,7 @@ __SCHEMA_TYPES = {
     'com.linkedin.pegasus2avro.retention.TimeBasedRetention': TimeBasedRetentionClass,
     'com.linkedin.pegasus2avro.retention.VersionBasedRetention': VersionBasedRetentionClass,
     'com.linkedin.pegasus2avro.role.Actors': ActorsClass,
+    'com.linkedin.pegasus2avro.role.RoleGroup': RoleGroupClass,
     'com.linkedin.pegasus2avro.role.RoleProperties': RolePropertiesClass,
     'com.linkedin.pegasus2avro.role.RoleUser': RoleUserClass,
     'com.linkedin.pegasus2avro.schema.ArrayType': ArrayTypeClass,
@@ -28067,6 +28106,7 @@ __SCHEMA_TYPES = {
     'TimeBasedRetention': TimeBasedRetentionClass,
     'VersionBasedRetention': VersionBasedRetentionClass,
     'Actors': ActorsClass,
+    'RoleGroup': RoleGroupClass,
     'RoleProperties': RolePropertiesClass,
     'RoleUser': RoleUserClass,
     'ArrayType': ArrayTypeClass,

datahub/metadata/com/linkedin/pegasus2avro/role/__init__.py

@@ -8,11 +8,13 @@
 # fmt: off
 # isort: skip_file
 from .....schema_classes import ActorsClass
+from .....schema_classes import RoleGroupClass
 from .....schema_classes import RolePropertiesClass
 from .....schema_classes import RoleUserClass
 
 
 Actors = ActorsClass
+RoleGroup = RoleGroupClass
 RoleProperties = RolePropertiesClass
 RoleUser = RoleUserClass
 

datahub/metadata/schema.avsc

@@ -3734,9 +3734,46 @@
         "name": "users",
         "default": null,
         "doc": "List of provisioned users of a role"
+      },
+      {
+        "type": [
+          "null",
+          {
+            "type": "array",
+            "items": {
+              "type": "record",
+              "name": "RoleGroup",
+              "namespace": "com.linkedin.pegasus2avro.role",
+              "fields": [
+                {
+                  "Relationship": {
+                    "entityTypes": [
+                      "corpGroup"
+                    ],
+                    "name": "Has"
+                  },
+                  "java": {
+                    "class": "com.linkedin.pegasus2avro.common.urn.Urn"
+                  },
+                  "Urn": "Urn",
+                  "entityTypes": [
+                    "corpGroup"
+                  ],
+                  "type": "string",
+                  "name": "group",
+                  "doc": "Link provisioned corp group for a role"
+                }
+              ],
+              "doc": "Provisioned groups of a role"
+            }
+          }
+        ],
+        "name": "groups",
+        "default": null,
+        "doc": "List of provisioned groups of a role"
       }
     ],
-    "doc": "Provisioned users of a role"
+    "doc": "Provisioned users and groups of a role"
   },
   {
     "type": "record",

datahub/metadata/schemas/Actors.avsc

@@ -42,7 +42,44 @@
       "name": "users",
       "default": null,
       "doc": "List of provisioned users of a role"
+    },
+    {
+      "type": [
+        "null",
+        {
+          "type": "array",
+          "items": {
+            "type": "record",
+            "name": "RoleGroup",
+            "namespace": "com.linkedin.pegasus2avro.role",
+            "fields": [
+              {
+                "Relationship": {
+                  "entityTypes": [
+                    "corpGroup"
+                  ],
+                  "name": "Has"
+                },
+                "java": {
+                  "class": "com.linkedin.pegasus2avro.common.urn.Urn"
+                },
+                "type": "string",
+                "name": "group",
+                "doc": "Link provisioned corp group for a role",
+                "Urn": "Urn",
+                "entityTypes": [
+                  "corpGroup"
+                ]
+              }
+            ],
+            "doc": "Provisioned groups of a role"
+          }
+        }
+      ],
+      "name": "groups",
+      "default": null,
+      "doc": "List of provisioned groups of a role"
     }
   ],
-  "doc": "Provisioned users of a role"
+  "doc": "Provisioned users and groups of a role"
 }

datahub/sdk/_shared.py

@@ -179,7 +179,7 @@ OwnerInputType: TypeAlias = Union[
     Tuple[ActorUrn, OwnershipTypeType],
     models.OwnerClass,
 ]
-OwnersInputType: TypeAlias = List[OwnerInputType]
+OwnersInputType: TypeAlias = Sequence[OwnerInputType]
 
 
 class HasOwnership(Entity):
@@ -280,7 +280,9 @@ class HasOwnership(Entity):
 # If you pass in a ContainerKey, we can use parent_key() to build the browse path.
 # If you pass in a list of urns, we'll use that as the browse path. Any non-urn strings
 # will be treated as raw ids.
-ParentContainerInputType: TypeAlias = Union["Container", ContainerKey, List[UrnOrStr]]
+ParentContainerInputType: TypeAlias = Union[
+    "Container", ContainerKey, Sequence[UrnOrStr]
+]
 
 
 class HasContainer(Entity):
@@ -340,7 +342,7 @@ class HasContainer(Entity):
                 )
                 for entry in parsed_path
             ]
-        elif container is not None:
+        elif isinstance(container, ContainerKey):
             container_urn = container.as_urn()
 
             browse_path_reversed = [container_urn]
@@ -399,7 +401,7 @@ class HasContainer(Entity):
 
 
 TagInputType: TypeAlias = Union[str, TagUrn, models.TagAssociationClass]
-TagsInputType: TypeAlias = List[TagInputType]
+TagsInputType: TypeAlias = Sequence[TagInputType]
 
 
 class HasTags(Entity):
@@ -454,7 +456,7 @@ class HasTags(Entity):
 TermInputType: TypeAlias = Union[
     str, GlossaryTermUrn, models.GlossaryTermAssociationClass
 ]
-TermsInputType: TypeAlias = List[TermInputType]
+TermsInputType: TypeAlias = Sequence[TermInputType]
 
 
 class HasTerms(Entity):

datahub/sdk/chart.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from datetime import datetime
-from typing import Dict, List, Optional, Type, Union
+from typing import Dict, List, Optional, Sequence, Type, Union
 
 from deprecated.sphinx import deprecated
 from typing_extensions import Self
@@ -73,7 +73,7 @@ class Chart(
         last_refreshed: Optional[datetime] = None,
         chart_type: Optional[Union[str, models.ChartTypeClass]] = None,
         access: Optional[str] = None,
-        input_datasets: Optional[List[Union[DatasetUrnOrStr, Dataset]]] = None,
+        input_datasets: Optional[Sequence[Union[DatasetUrnOrStr, Dataset]]] = None,
         # Standard aspects.
         parent_container: ParentContainerInputType | Unset = unset,
         subtype: Optional[str] = None,
@@ -291,7 +291,7 @@ class Chart(
         return [DatasetUrn.from_string(input_urn) for input_urn in (props.inputs or [])]
 
     def set_input_datasets(
-        self, input_datasets: List[Union[DatasetUrnOrStr, Dataset]]
+        self, input_datasets: Sequence[Union[DatasetUrnOrStr, Dataset]]
     ) -> None:
         """Set the input datasets of the chart."""
         # Convert all inputs to strings