feat(ingest): support full urns without owner_type in meta mapping (#…

…11298) Co-authored-by: Tamas Nemeth <[email protected]>
datahub-project · Sep 9, 2024 · 28310bb · 28310bb
1 parent 55a56be
commit 28310bb
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 25 deletions.
diff --git a/metadata-ingestion/src/datahub/utilities/mapping.py b/metadata-ingestion/src/datahub/utilities/mapping.py
@@ -383,14 +383,24 @@ def get_operation_value(
             if self.tag_prefix:
                 tag = self.tag_prefix + tag
             return tag
-        elif (
-            operation_type == Constants.ADD_OWNER_OPERATION
-            and operation_config[Constants.OWNER_TYPE]
-        ):
+        elif operation_type == Constants.ADD_OWNER_OPERATION:
             owner_id = _get_best_match(match, "owner")
-
             owner_ids: List[str] = [_id.strip() for _id in owner_id.split(",")]
 
+            owner_type_raw = operation_config.get(
+                Constants.OWNER_TYPE, Constants.USER_OWNER
+            )
+            owner_type_mapping: Dict[str, OwnerType] = {
+                Constants.USER_OWNER: OwnerType.USER,
+                Constants.GROUP_OWNER: OwnerType.GROUP,
+            }
+            if owner_type_raw not in owner_type_mapping:
+                logger.warning(
+                    f"Invalid owner type: {owner_type_raw}. Valid owner types are {', '.join(owner_type_mapping.keys())}"
+                )
+                return None
+            owner_type = owner_type_mapping[owner_type_raw]
+
             owner_category = (
                 operation_config.get(Constants.OWNER_CATEGORY)
                 or OwnershipTypeClass.DATAOWNER
@@ -402,19 +412,12 @@ def get_operation_value(
                     self.sanitize_owner_ids(owner_id) for owner_id in owner_ids
                 ]
 
-            owner_type_mapping: Dict[str, OwnerType] = {
-                Constants.USER_OWNER: OwnerType.USER,
-                Constants.GROUP_OWNER: OwnerType.GROUP,
-            }
-            if operation_config[Constants.OWNER_TYPE] in owner_type_mapping:
-                return _make_owner_category_list(
-                    owner_ids=owner_ids,
-                    owner_category=owner_category,
-                    owner_category_urn=owner_category_urn,
-                    owner_type=owner_type_mapping[
-                        operation_config[Constants.OWNER_TYPE]
-                    ],
-                )
+            return _make_owner_category_list(
+                owner_ids=owner_ids,
+                owner_category=owner_category,
+                owner_category_urn=owner_category_urn,
+                owner_type=owner_type,
+            )
 
         elif (
             operation_type == Constants.ADD_TERM_OPERATION

diff --git a/metadata-ingestion/tests/unit/test_mapping.py b/metadata-ingestion/tests/unit/test_mapping.py
@@ -188,6 +188,7 @@ def test_operation_processor_ownership_category():
         "user_owner": "@test_user",
         "business_owner": "alice,urn:li:corpGroup:biz-data-team",
         "architect": "bob",
+        "producer": "urn:li:corpGroup:producer-group",
     }
     processor = OperationProcessor(
         operation_defs={
@@ -215,33 +216,45 @@ def test_operation_processor_ownership_category():
                     "owner_category": "urn:li:ownershipType:architect",
                 },
             },
+            "producer": {
+                "match": ".*",
+                "operation": "add_owner",
+                "config": {
+                    # Testing using full urns without any owner_type set.
+                    "owner_category": OwnershipTypeClass.PRODUCER,
+                },
+            },
         },
         owner_source_type="SOURCE_CONTROL",
     )
     aspect_map = processor.process(raw_props)
     assert "add_owner" in aspect_map
 
     ownership_aspect: OwnershipClass = aspect_map["add_owner"]
-    assert len(ownership_aspect.owners) == 4
+    assert len(ownership_aspect.owners) == 5
+    assert all(
+        new_owner.source and new_owner.source.type == "SOURCE_CONTROL"
+        for new_owner in ownership_aspect.owners
+    )
 
     new_owner: OwnerClass = ownership_aspect.owners[0]
     assert new_owner.owner == "urn:li:corpGroup:biz-data-team"
-    assert new_owner.source and new_owner.source.type == "SOURCE_CONTROL"
     assert new_owner.type and new_owner.type == OwnershipTypeClass.BUSINESS_OWNER
 
     new_owner = ownership_aspect.owners[1]
+    assert new_owner.owner == "urn:li:corpGroup:producer-group"
+    assert new_owner.type and new_owner.type == OwnershipTypeClass.PRODUCER
+
+    new_owner = ownership_aspect.owners[2]
     assert new_owner.owner == "urn:li:corpGroup:test_user"
-    assert new_owner.source and new_owner.source.type == "SOURCE_CONTROL"
     assert new_owner.type and new_owner.type == OwnershipTypeClass.DATA_STEWARD
 
-    new_owner = ownership_aspect.owners[2]
+    new_owner = ownership_aspect.owners[3]
     assert new_owner.owner == "urn:li:corpuser:alice"
-    assert new_owner.source and new_owner.source.type == "SOURCE_CONTROL"
     assert new_owner.type and new_owner.type == OwnershipTypeClass.BUSINESS_OWNER
 
-    new_owner = ownership_aspect.owners[3]
+    new_owner = ownership_aspect.owners[4]
     assert new_owner.owner == "urn:li:corpuser:bob"
-    assert new_owner.source and new_owner.source.type == "SOURCE_CONTROL"
     assert new_owner.type == OwnershipTypeClass.CUSTOM
     assert new_owner.typeUrn == "urn:li:ownershipType:architect"