feat(ingest/dbt): support custom ownership types in dbt meta (#9332)

datahub-project · Dec 1, 2023 · 4d9eb12 · 4d9eb12
1 parent f9fd946
commit 4d9eb12
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 4 deletions.
diff --git a/metadata-ingestion/docs/sources/dbt/dbt.md b/metadata-ingestion/docs/sources/dbt/dbt.md
@@ -62,7 +62,7 @@ We support the following operations:
 1. add_tag - Requires `tag` property in config.
 2. add_term - Requires `term` property in config.
 3. add_terms - Accepts an optional `separator` property in config.
-4. add_owner - Requires `owner_type` property in config which can be either user or group. Optionally accepts the `owner_category` config property which you can set to one of `['TECHNICAL_OWNER', 'BUSINESS_OWNER', 'DATA_STEWARD', 'DATAOWNER'` (defaults to `DATAOWNER`).
+4. add_owner - Requires `owner_type` property in config which can be either user or group. Optionally accepts the `owner_category` config property which can be set to either a [custom ownership type](../../../../docs/ownership/ownership-types.md) urn like `urn:li:ownershipType:architect` or one of `['TECHNICAL_OWNER', 'BUSINESS_OWNER', 'DATA_STEWARD', 'DATAOWNER'` (defaults to `DATAOWNER`).
 5. add_doc_link - Requires `link` and `description` properties in config. Upon ingestion run, this will overwrite current links in the institutional knowledge section with this new link. The anchor text is defined here in the meta_mappings as `description`.
 
 Note:

diff --git a/metadata-ingestion/src/datahub/utilities/mapping.py b/metadata-ingestion/src/datahub/utilities/mapping.py
@@ -191,6 +191,7 @@ def convert_to_aspects(
                     OwnerClass(
                         owner=x.get("urn"),
                         type=x.get("category"),
+                        typeUrn=x.get("categoryUrn"),
                         source=OwnershipSourceClass(type=self.owner_source_type)
                         if self.owner_source_type
                         else None,
@@ -281,18 +282,25 @@ def get_operation_value(
                 operation_config.get(Constants.OWNER_CATEGORY)
                 or OwnershipTypeClass.DATAOWNER
             )
-            owner_category = owner_category.upper()
+            owner_category_urn = None
+            if owner_category.startswith("urn:li:"):
+                owner_category_urn = owner_category
+                owner_category = OwnershipTypeClass.DATAOWNER
+            else:
+                owner_category = owner_category.upper()
             if self.strip_owner_email_id:
                 owner_id = self.sanitize_owner_ids(owner_id)
             if operation_config[Constants.OWNER_TYPE] == Constants.USER_OWNER:
                 return {
                     "urn": mce_builder.make_owner_urn(owner_id, OwnerType.USER),
                     "category": owner_category,
+                    "categoryUrn": owner_category_urn,
                 }
             elif operation_config[Constants.OWNER_TYPE] == Constants.GROUP_OWNER:
                 return {
                     "urn": mce_builder.make_owner_urn(owner_id, OwnerType.GROUP),
                     "category": owner_category,
+                    "categoryUrn": owner_category_urn,
                 }
         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
@@ -174,7 +174,11 @@ def test_operation_processor_advanced_matching_owners():
 
 
 def test_operation_processor_ownership_category():
-    raw_props = {"user_owner": "@test_user", "business_owner": "alice"}
+    raw_props = {
+        "user_owner": "@test_user",
+        "business_owner": "alice",
+        "architect": "bob",
+    }
     processor = OperationProcessor(
         operation_defs={
             "user_owner": {
@@ -193,14 +197,22 @@ def test_operation_processor_ownership_category():
                     "owner_category": OwnershipTypeClass.BUSINESS_OWNER,
                 },
             },
+            "architect": {
+                "match": ".*",
+                "operation": "add_owner",
+                "config": {
+                    "owner_type": "user",
+                    "owner_category": "urn:li:ownershipType:architect",
+                },
+            },
         },
         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) == 2
+    assert len(ownership_aspect.owners) == 3
     new_owner: OwnerClass = ownership_aspect.owners[0]
     assert new_owner.owner == "urn:li:corpGroup:test_user"
     assert new_owner.source and new_owner.source.type == "SOURCE_CONTROL"
@@ -211,6 +223,12 @@ def test_operation_processor_ownership_category():
     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[2]
+    assert new_owner.owner == "urn:li:corpuser:bob"
+    assert new_owner.source and new_owner.source.type == "SOURCE_CONTROL"
+    assert new_owner.type == OwnershipTypeClass.DATAOWNER  # dummy value
+    assert new_owner.typeUrn == "urn:li:ownershipType:architect"
+
 
 def test_operation_processor_advanced_matching_tags():
     raw_props = {