Add OwnershipAssignment to dbt VirtualView (#653)

* Add `OwnershipAssignment` to dbt `VirtualView` * add config for owner assignment target * add quotes to ownership assignment choices description * parse string value with commas as sequence of strings * fix readme example * bump version * fix readme * bump version * Revert "parse string value with commas as sequence of strings" This reverts commit f7ca70f. * bump version --------- Signed-off-by: Tsung-Ju Lii <[email protected]>
MetaphorData · Nov 1, 2023 · 63375e0 · 63375e0
1 parent 3733d77
commit 63375e0
Show file tree

Hide file tree

Showing 12 changed files with 157 additions and 14 deletions.
diff --git a/metaphor/dbt/README.md b/metaphor/dbt/README.md
@@ -72,6 +72,15 @@ meta_ownerships:
     email_domain: test.com
 ```
 
+You can also choose between assigning the owner to the materialized table, the dbt model, or both:
+
+```yaml
+meta_ownerships:
+  - meta_key: owner
+    ownership_type: Data Steward
+    assignment_target: dbt_model # Valid choices: "dbt_model", "materialized_table", "both". Default is "both"
+```
+
 #### Governed Tags
 
 Similar to [Ownership](#ownership), you can optionally specify certain attributes in meta. For example:

diff --git a/metaphor/dbt/config.py b/metaphor/dbt/config.py
@@ -1,11 +1,13 @@
 from dataclasses import field as dataclass_field
-from typing import List, Optional
+from typing import List, Literal, Optional
 
 from pydantic.dataclasses import dataclass
 
 from metaphor.common.base_config import BaseConfig
 from metaphor.common.dataclass import ConnectorConfig
 
+MetaOwnershipAssignmentTarget = Literal["dbt_model", "materialized_table", "both"]
+
 
 @dataclass(config=ConnectorConfig)
 class MetaOwnership:
@@ -18,6 +20,10 @@ class MetaOwnership:
     # Domain for user names
     email_domain: Optional[str] = None
 
+    # The target to assign this ownership to. Can be either the dbt model, the materialized table, or both.
+    # Defaults to both.
+    assignment_target: MetaOwnershipAssignmentTarget = "both"
+
 
 @dataclass(config=ConnectorConfig)
 class MetaTag:

diff --git a/metaphor/dbt/manifest_parser.py b/metaphor/dbt/manifest_parser.py
@@ -418,7 +418,7 @@ def _parse_model(
             virtual_view.dbt_model.compiled_sql = getattr(model, "compiled_code")
             dbt_model.compiled_sql = getattr(model, "compiled_code")
 
-        self._parse_model_meta(model)
+        self._parse_model_meta(model, virtual_view)
 
         self._parse_model_materialization(model, dbt_model)
 
@@ -458,7 +458,9 @@ def _parse_macros(self, macros: MACRO_MAP) -> Dict[str, DbtMacro]:
 
         return macro_map
 
-    def _parse_model_meta(self, model: MODEL_NODE_TYPE) -> None:
+    def _parse_model_meta(
+        self, model: MODEL_NODE_TYPE, virtual_view: VirtualView
+    ) -> None:
         if model.config is None or model.database is None:
             logger.warning("Skipping model without config or database")
             return
@@ -488,9 +490,13 @@ def get_dataset():
 
         # Assign ownership & tags to materialized table/view
         ownerships = get_ownerships_from_meta(meta, self._meta_ownerships)
-        if len(ownerships) > 0:
+        if len(ownerships.materialized_table) > 0:
             get_dataset().ownership_assignment = OwnershipAssignment(
-                ownerships=ownerships
+                ownerships=ownerships.materialized_table
+            )
+        if len(ownerships.dbt_model) > 0:
+            virtual_view.ownership_assignment = OwnershipAssignment(
+                ownerships=ownerships.dbt_model
             )
 
         tag_names = get_tags_from_meta(meta, self._meta_tags)

diff --git a/metaphor/dbt/util.py b/metaphor/dbt/util.py
@@ -1,4 +1,5 @@
 import re
+from dataclasses import dataclass, field
 from typing import Dict, List, Optional
 
 from metaphor.common.entity_id import (
@@ -55,9 +56,15 @@ def get_metric_name_from_unique_id(unique_id: str) -> str:
     return unique_id[7:]
 
 
+@dataclass
+class Ownerships:
+    dbt_model: List[Ownership] = field(default_factory=lambda: [])
+    materialized_table: List[Ownership] = field(default_factory=lambda: [])
+
+
 def get_ownerships_from_meta(
     meta: Dict, meta_ownerships: List[MetaOwnership]
-) -> List[Ownership]:
+) -> Ownerships:
     """Extracts ownership info from models' meta field"""
 
     def to_owner(email_or_username: str, email_domain: Optional[str]) -> Optional[str]:
@@ -71,7 +78,7 @@ def to_owner(email_or_username: str, email_domain: Optional[str]) -> Optional[st
 
         return str(to_person_entity_id(email))
 
-    ownerships: List[Ownership] = []
+    ownerships = Ownerships()
     for meta_ownership in meta_ownerships:
         value = meta.get(meta_ownership.meta_key)
         email_or_usernames = []
@@ -83,12 +90,14 @@ def to_owner(email_or_username: str, email_domain: Optional[str]) -> Optional[st
         for email_or_username in email_or_usernames:
             owner = to_owner(email_or_username, meta_ownership.email_domain)
             if owner is not None:
-                ownerships.append(
-                    Ownership(
-                        contact_designation_name=meta_ownership.ownership_type,
-                        person=owner,
-                    )
+                ownership = Ownership(
+                    contact_designation_name=meta_ownership.ownership_type,
+                    person=owner,
                 )
+                if meta_ownership.assignment_target in ["dbt_model", "both"]:
+                    ownerships.dbt_model.append(ownership)
+                if meta_ownership.assignment_target in ["materialized_table", "both"]:
+                    ownerships.materialized_table.append(ownership)
 
     return ownerships
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "metaphor-connectors"
-version = "0.13.27"
+version = "0.13.28"
 license = "Apache-2.0"
 description = "A collection of Python-based 'connectors' that extract metadata from various sources to ingest into the Metaphor app."
 authors = ["Metaphor <[email protected]>"]

diff --git a/tests/dbt/data/trial_v4/results.json b/tests/dbt/data/trial_v4/results.json
@@ -211,6 +211,14 @@
     "logicalId": {
       "name": "trial.my_first_dbt_model",
       "type": "DBT_MODEL"
+    },
+    "ownershipAssignment": {
+      "ownerships": [
+        {
+          "contactDesignationName": "Maintainer",
+          "person": "PERSON~8B6974C32BFCBBF6AB3930E1B7A17846"
+        }
+      ]
     }
   },
   {

diff --git a/tests/dbt/data/trial_v5/results.json b/tests/dbt/data/trial_v5/results.json
@@ -211,6 +211,14 @@
     "logicalId": {
       "name": "trial.my_first_dbt_model",
       "type": "DBT_MODEL"
+    },
+    "ownershipAssignment": {
+      "ownerships": [
+        {
+          "contactDesignationName": "Maintainer",
+          "person": "PERSON~8B6974C32BFCBBF6AB3930E1B7A17846"
+        }
+      ]
     }
   },
   {

diff --git a/tests/dbt/data/trial_v6/results.json b/tests/dbt/data/trial_v6/results.json
@@ -211,6 +211,14 @@
     "logicalId": {
       "name": "trial.my_first_dbt_model",
       "type": "DBT_MODEL"
+    },
+    "ownershipAssignment": {
+      "ownerships": [
+        {
+          "contactDesignationName": "Maintainer",
+          "person": "PERSON~8B6974C32BFCBBF6AB3930E1B7A17846"
+        }
+      ]
     }
   },
   {

diff --git a/tests/dbt/data/trial_v7/results.json b/tests/dbt/data/trial_v7/results.json
@@ -211,6 +211,14 @@
     "logicalId": {
       "name": "trial.my_first_dbt_model",
       "type": "DBT_MODEL"
+    },
+    "ownershipAssignment": {
+      "ownerships": [
+        {
+          "contactDesignationName": "Maintainer",
+          "person": "PERSON~8B6974C32BFCBBF6AB3930E1B7A17846"
+        }
+      ]
     }
   },
   {

diff --git a/tests/dbt/data/trial_v8/results.json b/tests/dbt/data/trial_v8/results.json
@@ -211,6 +211,14 @@
     "logicalId": {
       "name": "trial.my_first_dbt_model",
       "type": "DBT_MODEL"
+    },
+    "ownershipAssignment": {
+      "ownerships": [
+        {
+          "contactDesignationName": "Maintainer",
+          "person": "PERSON~8B6974C32BFCBBF6AB3930E1B7A17846"
+        }
+      ]
     }
   },
   {

diff --git a/tests/dbt/data/trial_v9/results.json b/tests/dbt/data/trial_v9/results.json
@@ -211,6 +211,14 @@
     "logicalId": {
       "name": "trial.my_first_dbt_model",
       "type": "DBT_MODEL"
+    },
+    "ownershipAssignment": {
+      "ownerships": [
+        {
+          "contactDesignationName": "Maintainer",
+          "person": "PERSON~8B6974C32BFCBBF6AB3930E1B7A17846"
+        }
+      ]
     }
   },
   {

diff --git a/tests/dbt/test_util.py b/tests/dbt/test_util.py
@@ -46,7 +46,72 @@ def test_get_ownerships_from_meta(test_root_dir):
         ),
     ]
 
-    assert get_ownerships_from_meta(meta, meta_ownerships) == expected_ownerships
+    assert (
+        get_ownerships_from_meta(meta, meta_ownerships).materialized_table
+        == expected_ownerships
+    )
+    assert (
+        get_ownerships_from_meta(meta, meta_ownerships).dbt_model == expected_ownerships
+    )
+
+
+def test_get_ownerships_with_assignment_targets(test_root_dir):
+    meta = {
+        "owners_dbt_model": ["foo", "bar"],
+        "owners_materialized_table": ["bar", "qux"],
+        "owners_both": ["baz"],
+    }
+    meta_ownerships = [
+        MetaOwnership(
+            meta_key="owners_dbt_model",
+            ownership_type="dbt model owner",
+            email_domain="metaphor.io",
+            assignment_target="dbt_model",
+        ),
+        MetaOwnership(
+            meta_key="owners_both",
+            ownership_type="owner of both dbt model and materialized table",
+            email_domain="metaphor.io",
+            assignment_target="both",
+        ),
+        MetaOwnership(
+            meta_key="owners_materialized_table",
+            ownership_type="materialized table owner",
+            email_domain="metaphor.io",
+            assignment_target="materialized_table",
+        ),
+    ]
+    ownerships = get_ownerships_from_meta(meta, meta_ownerships)
+    expected_dbt_model_ownerships = [
+        Ownership(
+            contact_designation_name="dbt model owner",
+            person=str(to_person_entity_id("[email protected]")),
+        ),
+        Ownership(
+            contact_designation_name="dbt model owner",
+            person=str(to_person_entity_id("[email protected]")),
+        ),
+        Ownership(
+            contact_designation_name="owner of both dbt model and materialized table",
+            person=str(to_person_entity_id("[email protected]")),
+        ),
+    ]
+    expected_materialized_table_ownerships = [
+        Ownership(
+            contact_designation_name="owner of both dbt model and materialized table",
+            person=str(to_person_entity_id("[email protected]")),
+        ),
+        Ownership(
+            contact_designation_name="materialized table owner",
+            person=str(to_person_entity_id("[email protected]")),
+        ),
+        Ownership(
+            contact_designation_name="materialized table owner",
+            person=str(to_person_entity_id("[email protected]")),
+        ),
+    ]
+    assert ownerships.dbt_model == expected_dbt_model_ownerships
+    assert ownerships.materialized_table == expected_materialized_table_ownerships
 
 
 def test_get_tags_from_meta(test_root_dir):