Resolve service principal owners to their display names

metaphor/unity_catalog/extractor.py

@@ -6,6 +6,7 @@
 from typing import Collection, Dict, Generator, Iterator, List, Optional, Tuple
 
 from databricks.sdk.service.catalog import TableInfo, TableType, VolumeInfo
+from databricks.sdk.service.iam import ServicePrincipal
 
 from metaphor.common.base_extractor import BaseExtractor
 from metaphor.common.entity_id import (
@@ -62,6 +63,7 @@
     create_connection_pool,
     list_column_lineage,
     list_query_log,
+    list_service_principals,
     list_table_lineage,
 )
 
@@ -164,6 +166,8 @@ def __init__(self, config: UnityCatalogRunConfig):
             http_path=config.http_path,
         )
 
+        self._service_principals: Dict[str, ServicePrincipal] = {}
+
         self._last_refresh_time_queue: List[str] = []
 
         # Map fullname or volume path to a dataset
@@ -176,6 +180,9 @@ def __init__(self, config: UnityCatalogRunConfig):
     async def extract(self) -> Collection[ENTITY_TYPES]:
         logger.info("Fetching metadata from Unity Catalog")
 
+        self._service_principals = list_service_principals(self._api)
+        logger.info(f"Found service principals: {self._service_principals}")
+
         catalogs = [
             catalog
             for catalog in self._get_catalogs()
@@ -347,10 +354,19 @@ def _init_dataset(self, table_info: TableInfo) -> Dataset:
         )
 
         if table_info.owner is not None:
+            # Unity Catalog returns service principal's application_id and must be
+            # manually map back to display_name
+            service_principal = self._service_principals.get(table_info.owner)
+            owner = (
+                service_principal.display_name
+                if service_principal
+                else table_info.owner
+            )
+
             dataset.system_contacts = SystemContacts(
                 contacts=[
                     SystemContact(
-                        email=table_info.owner,
+                        email=owner,
                         system_contact_source=SystemContactSource.UNITY_CATALOG,
                     )
                 ]

metaphor/unity_catalog/utils.py

@@ -5,6 +5,7 @@
 
 from databricks import sql
 from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.iam import ServicePrincipal
 from databricks.sql.client import Connection
 
 from metaphor.common.logger import get_logger, json_dump_to_debug_file
@@ -121,6 +122,7 @@
     connection: Connection, lookback_days: int, excluded_usernames: Collection[str]
 ):
     """
+    Fetch query logs from system.query.history table
     See https://docs.databricks.com/en/admin/system-tables/query-history.html
     """
     start = start_of_day(lookback_days)
@@ -268,3 +270,29 @@
 
 def create_api(host: str, token: str) -> WorkspaceClient:
     return WorkspaceClient(host=host, token=token)
+
+
+def list_service_principals(api: WorkspaceClient) -> Dict[str, ServicePrincipal]:
+    """
+    List all service principals
+    See https://docs.databricks.com/api/workspace/groups/list
+    """
+
+    # See https://databricks-sdk-py.readthedocs.io/en/latest/dbdataclasses/iam.html#databricks.sdk.service.iam.Group
+    excluded_attributes = "entitlements,groups,roles,schemas"
+
+    try:
+        principals = list(
+            api.service_principals.list(excluded_attributes=excluded_attributes)
+        )
+        json_dump_to_debug_file(principals, "list-principals.json")
+
+        principal_map: Dict[str, ServicePrincipal] = {}
+        for p in principals:
+            if p.application_id is not None:
+                principal_map[p.application_id] = p
+
+        return principal_map
+    except Exception as e:
+        logger.error(f"Failed to list principals: {e}")
+        return {}

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "metaphor-connectors"
-version = "0.14.48"
+version = "0.14.49"
 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]>"]

tests/unity_catalog/expected.json

@@ -374,7 +374,7 @@
     "systemContacts": {
       "contacts": [
         {
-          "email": "[email protected]",
+          "email": "service principal 1",
           "systemContactSource": "UNITY_CATALOG"
         }
       ]
@@ -386,7 +386,7 @@
       "datasetType": "UNITY_CATALOG_TABLE",
       "tableInfo": {
         "dataSourceFormat": "DELTA",
-        "owner": "[email protected]",
+        "owner": "sp1",
         "properties": [
           {
             "key": "delta.lastCommitTimestamp",

tests/unity_catalog/test_extractor.py

@@ -12,6 +12,7 @@
 )
 from databricks.sdk.service.catalog import TableInfo as Table
 from databricks.sdk.service.catalog import TableType, VolumeInfo, VolumeType
+from databricks.sdk.service.iam import ServicePrincipal
 from pytest_snapshot.plugin import Snapshot
 
 from metaphor.common.base_config import OutputConfig
@@ -39,8 +40,10 @@ def dummy_config():
 @patch("metaphor.unity_catalog.extractor.list_table_lineage")
 @patch("metaphor.unity_catalog.extractor.list_column_lineage")
 @patch("metaphor.unity_catalog.extractor.list_query_log")
+@patch("metaphor.unity_catalog.extractor.list_service_principals")
 @pytest.mark.asyncio
 async def test_extractor(
+    mock_list_service_principals: MagicMock,
     mock_list_query_log: MagicMock,
     mock_list_column_lineage: MagicMock,
     mock_list_table_lineage: MagicMock,
@@ -50,6 +53,9 @@ async def test_extractor(
     mock_batch_get_last_refreshed_time: MagicMock,
     test_root_dir: str,
 ):
+    mock_list_service_principals.return_value = {
+        "sp1": ServicePrincipal(display_name="service principal 1")
+    }
 
     def mock_list_catalogs():
         return [CatalogInfo(name="catalog")]
@@ -137,7 +143,7 @@ def mock_list_tables(catalog, schema):
                     )
                 ],
                 storage_location="s3://path",
-                owner="[email protected]",
+                owner="sp1",
                 comment="example",
                 updated_at=0,
                 updated_by="[email protected]",

tests/unity_catalog/test_utils.py

@@ -1,6 +1,7 @@
 from datetime import datetime, timezone
 from unittest.mock import MagicMock
 
+from databricks.sdk.service.iam import ServicePrincipal
 from freezegun import freeze_time
 from pytest_snapshot.plugin import Snapshot
 
@@ -11,6 +12,7 @@
     get_last_refreshed_time,
     list_column_lineage,
     list_query_log,
+    list_service_principals,
     list_table_lineage,
 )
 from tests.unity_catalog.mocks import mock_connection_pool, mock_sql_connection
@@ -152,6 +154,20 @@ def test_batch_get_last_refreshed_time():
     assert result_map == {"a.b.c": datetime(2020, 1, 1), "d.e.f": datetime(2020, 1, 1)}
 
 
+def test_list_service_principals():
+
+    sp1 = ServicePrincipal(application_id="sp1", display_name="SP1")
+    sp2 = ServicePrincipal(application_id="sp2", display_name="SP2")
+
+    mock_api = MagicMock()
+    mock_api.service_principals = MagicMock()
+    mock_api.service_principals.list.return_value = [sp1, sp2]
+
+    principals = list_service_principals(mock_api)
+
+    assert principals == {"sp1": sp1, "sp2": sp2}
+
+
 def test_escape_special_characters():
     assert escape_special_characters("this.is.a_table") == "this.is.a_table"
     assert escape_special_characters("this.is.also-a-table") == "`this.is.also-a-table`"