feat(ingestion): Add typeUrn handling to ownership transformers

metadata-ingestion/docs/transformer/dataset_transformer.md

@@ -55,12 +55,12 @@ transformers:
 ```
 ## Simple Add Dataset ownership 
 ### Config Details
-| Field                       | Required | Type         | Default       | Description                                                      |
-|-----------------------------|----------|--------------|---------------|------------------------------------------------------------------|
-| `owner_urns`                | ✅        | list[string] |               | List of owner urns.                                              |
-| `ownership_type`            |          | string       | `DATAOWNER`   | ownership type of the owners.                                    |
-| `replace_existing`          |          | boolean      | `false`       | Whether to remove owners from entity sent by ingestion source.   |
-| `semantics`                 |          | enum         | `OVERWRITE`   | Whether to OVERWRITE or PATCH the entity present on DataHub GMS. |
+| Field              | Required | Type         | Default     | Description                                                         |
+|--------------------|----------|--------------|-------------|---------------------------------------------------------------------|
+| `owner_urns`       | ✅        | list[string] |             | List of owner urns.                                                 |
+| `ownership_type`   |          | string       | "DATAOWNER" | ownership type of the owners (either as enum or ownership type urn) |
+| `replace_existing` |          | boolean      | `false`     | Whether to remove owners from entity sent by ingestion source.      |
+| `semantics`        |          | enum         | `OVERWRITE` | Whether to OVERWRITE or PATCH the entity present on DataHub GMS.    |
 
 For transformer behaviour on `replace_existing` and `semantics`, please refer section [Relationship Between replace_existing And semantics](#relationship-between-replace_existing-and-semantics).
 
@@ -95,7 +95,7 @@ transformers:
             - "urn:li:corpuser:username1"
             - "urn:li:corpuser:username2"
             - "urn:li:corpGroup:groupname"
-          ownership_type: "PRODUCER"
+          ownership_type: "urn:li:ownershipType:__system__producer"
     ```
 - Add owners, however overwrite the owners available for the dataset on DataHub GMS
     ```yaml
@@ -107,7 +107,7 @@ transformers:
             - "urn:li:corpuser:username1"
             - "urn:li:corpuser:username2"
             - "urn:li:corpGroup:groupname"
-          ownership_type: "PRODUCER"
+          ownership_type: "urn:li:ownershipType:__system__producer" 
     ```
 - Add owners, however keep the owners available for the dataset on DataHub GMS
     ```yaml
@@ -124,12 +124,12 @@ transformers:
 
 ## Pattern Add Dataset ownership 
 ### Config Details
-| Field                       | Required | Type                  | Default          | Description                                                                             |
-|-----------------------------|--------- |-----------------------|------------------|-----------------------------------------------------------------------------------------|
-| `owner_pattern`             | ✅        | map[regx, list[urn]]  |                  | entity urn with regular expression and list of owners urn apply to matching entity urn. |
-| `ownership_type`            |          | string                | `DATAOWNER`      | ownership type of the owners.                                                           |
-| `replace_existing`          |          | boolean               | `false`          | Whether to remove owners from entity sent by ingestion source.                          |
-| `semantics`                 |          | enum                  | `OVERWRITE`      | Whether to OVERWRITE or PATCH the entity present on DataHub GMS.                        |
+| Field              | Required | Type                 | Default     | Description                                                                             |
+|--------------------|----------|----------------------|-------------|-----------------------------------------------------------------------------------------|
+| `owner_pattern`    | ✅        | map[regx, list[urn]] |             | entity urn with regular expression and list of owners urn apply to matching entity urn. |
+| `ownership_type`   |          | string               | "DATAOWNER" | ownership type of the owners (either as enum or ownership type urn)                     |
+| `replace_existing` |          | boolean              | `false`     | Whether to remove owners from entity sent by ingestion source.                          |
+| `semantics`        |          | enum                 | `OVERWRITE` | Whether to OVERWRITE or PATCH the entity present on DataHub GMS.                        |
 
 let’s suppose we’d like to append a series of users who we know to own a different dataset from a data source but aren't detected during normal ingestion. To do so, we can use the `pattern_add_dataset_ownership` module that’s included in the ingestion framework.  This will match the pattern to `urn` of the dataset and assign the respective owners.
 
@@ -158,7 +158,7 @@ The config, which we’d append to our ingestion recipe YAML, would look like th
             rules:
               ".*example1.*": ["urn:li:corpuser:username1"]
               ".*example2.*": ["urn:li:corpuser:username2"]
-          ownership_type: "PRODUCER"
+          ownership_type: "urn:li:ownershipType:__system__producer"
     ```
 - Add owner, however overwrite the owners available for the dataset on DataHub GMS
     ```yaml
@@ -170,7 +170,7 @@ The config, which we’d append to our ingestion recipe YAML, would look like th
             rules:
               ".*example1.*": ["urn:li:corpuser:username1"]
               ".*example2.*": ["urn:li:corpuser:username2"]
-          ownership_type: "PRODUCER"
+          ownership_type: "urn:li:ownershipType:__system__producer" 
     ```
 - Add owner, however keep the owners available for the dataset on DataHub GMS
     ```yaml

metadata-ingestion/src/datahub/emitter/mce_builder.py

@@ -9,12 +9,13 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    Iterable,
     List,
     Optional,
+    Tuple,
     Type,
     TypeVar,
     Union,
-    cast,
     get_type_hints,
 )
 
@@ -342,26 +343,20 @@ def make_ml_model_group_urn(platform: str, group_name: str, env: str) -> str:
     )
 
 
-def is_valid_ownership_type(ownership_type: Optional[str]) -> bool:
-    return ownership_type is not None and ownership_type in [
-        OwnershipTypeClass.TECHNICAL_OWNER,
-        OwnershipTypeClass.BUSINESS_OWNER,
-        OwnershipTypeClass.DATA_STEWARD,
-        OwnershipTypeClass.NONE,
-        OwnershipTypeClass.DEVELOPER,
-        OwnershipTypeClass.DATAOWNER,
-        OwnershipTypeClass.DELEGATE,
-        OwnershipTypeClass.PRODUCER,
-        OwnershipTypeClass.CONSUMER,
-        OwnershipTypeClass.STAKEHOLDER,
+def get_class_fields(_class: Type[object]) -> Iterable[str]:
+    return [
+        f
+        for f in dir(_class)
+        if not callable(getattr(_class, f)) and not f.startswith("_")
     ]
 
 
-def validate_ownership_type(ownership_type: Optional[str]) -> str:
-    if is_valid_ownership_type(ownership_type):
-        return cast(str, ownership_type)
-    else:
-        raise ValueError(f"Unexpected ownership type: {ownership_type}")
+def validate_ownership_type(ownership_type: str) -> Tuple[str, Optional[str]]:
+    if ownership_type.startswith("urn:li:"):
+        return OwnershipTypeClass.CUSTOM, ownership_type
+    if ownership_type in get_class_fields(OwnershipTypeClass):
+        return ownership_type, None
+    raise ValueError(f"Unexpected ownership type: {ownership_type}")
 
 
 def make_lineage_mce(

metadata-ingestion/src/datahub/ingestion/transformer/add_dataset_ownership.py

@@ -14,11 +14,8 @@
 from datahub.ingestion.transformer.dataset_transformer import (
     DatasetOwnershipTransformer,
 )
-from datahub.metadata.schema_classes import (
-    OwnerClass,
-    OwnershipClass,
-    OwnershipTypeClass,
-)
+from datahub.metadata._schema_classes import OwnershipTypeClass
+from datahub.metadata.schema_classes import OwnerClass, OwnershipClass
 
 
 class AddDatasetOwnershipConfig(TransformerSemanticsConfigModel):
@@ -102,7 +99,7 @@ def transform_aspect(
 
 
 class DatasetOwnershipBaseConfig(TransformerSemanticsConfigModel):
-    ownership_type: Optional[str] = OwnershipTypeClass.DATAOWNER
+    ownership_type: str = OwnershipTypeClass.DATAOWNER
 
 
 class SimpleDatasetOwnershipConfig(DatasetOwnershipBaseConfig):
@@ -114,11 +111,14 @@ class SimpleAddDatasetOwnership(AddDatasetOwnership):
     """Transformer that adds a specified set of owners to each dataset."""
 
     def __init__(self, config: SimpleDatasetOwnershipConfig, ctx: PipelineContext):
-        ownership_type = builder.validate_ownership_type(config.ownership_type)
+        ownership_type, ownership_type_urn = builder.validate_ownership_type(
+            config.ownership_type
+        )
         owners = [
             OwnerClass(
                 owner=owner,
                 type=ownership_type,
+                typeUrn=ownership_type_urn,
             )
             for owner in config.owner_urns
         ]
@@ -147,29 +147,17 @@ class PatternDatasetOwnershipConfig(DatasetOwnershipBaseConfig):
 class PatternAddDatasetOwnership(AddDatasetOwnership):
     """Transformer that adds a specified set of owners to each dataset."""
 
-    def getOwners(
-        self,
-        key: str,
-        owner_pattern: KeyValuePattern,
-        ownership_type: Optional[str] = None,
-    ) -> List[OwnerClass]:
-        owners = [
-            OwnerClass(
-                owner=owner,
-                type=builder.validate_ownership_type(ownership_type),
-            )
-            for owner in owner_pattern.value(key)
-        ]
-        return owners
-
     def __init__(self, config: PatternDatasetOwnershipConfig, ctx: PipelineContext):
-        ownership_type = builder.validate_ownership_type(config.ownership_type)
         owner_pattern = config.owner_pattern
+        ownership_type, ownership_type_urn = builder.validate_ownership_type(
+            config.ownership_type
+        )
         generic_config = AddDatasetOwnershipConfig(
             get_owners_to_add=lambda urn: [
                 OwnerClass(
                     owner=owner,
                     type=ownership_type,
+                    typeUrn=ownership_type_urn,
                 )
                 for owner in owner_pattern.value(urn)
             ],

metadata-ingestion/tests/unit/test_pipeline.py

@@ -214,7 +214,10 @@ def test_run_including_registered_transformation(self):
                 "transformers": [
                     {
                         "type": "simple_add_dataset_ownership",
-                        "config": {"owner_urns": ["urn:li:corpuser:foo"]},
+                        "config": {
+                            "owner_urns": ["urn:li:corpuser:foo"],
+                            "ownership_type": "urn:li:ownershipType:__system__technical_owner",
+                        },
                     }
                 ],
                 "sink": {"type": "tests.test_helpers.sink_helpers.RecordingSink"},

metadata-ingestion/tests/unit/test_transform_dataset.py

@@ -234,7 +234,7 @@ def test_simple_dataset_ownership_transformation(mock_time):
     assert last_event.entityUrn == outputs[0].record.proposedSnapshot.urn
     assert all(
         [
-            owner.type == models.OwnershipTypeClass.DATAOWNER
+            owner.type == models.OwnershipTypeClass.DATAOWNER and owner.typeUrn is None
             for owner in last_event.aspect.owners
         ]
     )
@@ -247,7 +247,7 @@ def test_simple_dataset_ownership_transformation(mock_time):
     assert len(second_ownership_aspect.owners) == 3
     assert all(
         [
-            owner.type == models.OwnershipTypeClass.DATAOWNER
+            owner.type == models.OwnershipTypeClass.DATAOWNER and owner.typeUrn is None
             for owner in second_ownership_aspect.owners
         ]
     )
@@ -293,6 +293,44 @@ def test_simple_dataset_ownership_with_type_transformation(mock_time):
     assert ownership_aspect.owners[0].type == models.OwnershipTypeClass.PRODUCER
 
 
+def test_simple_dataset_ownership_with_type_urn_transformation(mock_time):
+    input = make_generic_dataset()
+
+    transformer = SimpleAddDatasetOwnership.create(
+        {
+            "owner_urns": [
+                builder.make_user_urn("person1"),
+            ],
+            "ownership_type": "urn:li:ownershipType:__system__technical_owner",
+        },
+        PipelineContext(run_id="test"),
+    )
+
+    output = list(
+        transformer.transform(
+            [
+                RecordEnvelope(input, metadata={}),
+                RecordEnvelope(EndOfStream(), metadata={}),
+            ]
+        )
+    )
+
+    assert len(output) == 3
+
+    # original MCE is unchanged
+    assert input == output[0].record
+
+    ownership_aspect = output[1].record.aspect
+
+    assert isinstance(ownership_aspect, OwnershipClass)
+    assert len(ownership_aspect.owners) == 1
+    assert ownership_aspect.owners[0].type == OwnershipTypeClass.CUSTOM
+    assert (
+        ownership_aspect.owners[0].typeUrn
+        == "urn:li:ownershipType:__system__technical_owner"
+    )
+
+
 def _test_extract_tags(in_urn: str, regex_str: str, out_tag: str) -> None:
     input = make_generic_dataset(entity_urn=in_urn)
     transformer = ExtractDatasetTags.create(
@@ -883,6 +921,7 @@ def test_pattern_dataset_ownership_transformation(mock_time):
                     ".*example2.*": [builder.make_user_urn("person2")],
                 }
             },
+            "ownership_type": "DATAOWNER",
         },
         PipelineContext(run_id="test"),
     )
@@ -2233,6 +2272,7 @@ def fake_ownership_class(entity_urn: str) -> models.OwnershipClass:
             "replace_existing": False,
             "semantics": TransformerSemantics.PATCH,
             "owner_urns": [owner2],
+            "ownership_type": "DATAOWNER",
         },
         pipeline_context=pipeline_context,
     )