From 75a3a10bc17406630ec93dccd024d4e53c31cb8c Mon Sep 17 00:00:00 2001
From: Harshal Sheth <hsheth2@gmail.com>
Date: Wed, 21 Aug 2024 16:10:51 -0400
Subject: [PATCH] feat(ingest/dbt): add support for urns in add_owner directive

---
 metadata-ingestion/docs/sources/dbt/dbt.md         |  6 +++++-
 .../src/datahub/emitter/mce_builder.py             |  6 ++++++
 metadata-ingestion/tests/unit/test_mapping.py      | 14 ++++++++++----
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/metadata-ingestion/docs/sources/dbt/dbt.md b/metadata-ingestion/docs/sources/dbt/dbt.md
index eca5101e00642..2333ddcee677b 100644
--- a/metadata-ingestion/docs/sources/dbt/dbt.md
+++ b/metadata-ingestion/docs/sources/dbt/dbt.md
@@ -62,7 +62,11 @@ 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 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`).
+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`).
+
+   - The `owner_type` property will be ignored if the owner is a fully qualified urn.
+   - You can use commas to specify multiple owners - e.g. `business_owner: "jane,john,urn:li:corpGroup:data-team"`.
+
 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/emitter/mce_builder.py b/metadata-ingestion/src/datahub/emitter/mce_builder.py
index a0694983660e4..e273bab62fe7a 100644
--- a/metadata-ingestion/src/datahub/emitter/mce_builder.py
+++ b/metadata-ingestion/src/datahub/emitter/mce_builder.py
@@ -244,6 +244,12 @@ def make_tag_urn(tag: str) -> str:
 
 
 def make_owner_urn(owner: str, owner_type: OwnerType) -> str:
+    if owner_type == OwnerType.USER:
+        return make_user_urn(owner)
+    elif owner_type == OwnerType.GROUP:
+        return make_group_urn(owner)
+    # This should pretty much never happen.
+    # TODO: With Python 3.11, we can use typing.assert_never() here.
     return f"urn:li:{owner_type.value}:{owner}"
 
 
diff --git a/metadata-ingestion/tests/unit/test_mapping.py b/metadata-ingestion/tests/unit/test_mapping.py
index 0e176710bb4f3..b4168af3029d0 100644
--- a/metadata-ingestion/tests/unit/test_mapping.py
+++ b/metadata-ingestion/tests/unit/test_mapping.py
@@ -186,7 +186,7 @@ def test_operation_processor_advanced_matching_owners():
 def test_operation_processor_ownership_category():
     raw_props = {
         "user_owner": "@test_user",
-        "business_owner": "alice",
+        "business_owner": "alice,urn:li:corpGroup:biz-data-team",
         "architect": "bob",
     }
     processor = OperationProcessor(
@@ -222,18 +222,24 @@ def test_operation_processor_ownership_category():
     assert "add_owner" in aspect_map
 
     ownership_aspect: OwnershipClass = aspect_map["add_owner"]
-    assert len(ownership_aspect.owners) == 3
+    assert len(ownership_aspect.owners) == 4
+
     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: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[1]
+    new_owner = ownership_aspect.owners[2]
     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[2]
+    new_owner = ownership_aspect.owners[3]
     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