Implementation of an Ownership factory

src/databricks/labs/ucx/assessment/clusters.py

@@ -30,7 +30,7 @@
 )
 from databricks.labs.ucx.assessment.init_scripts import CheckInitScriptMixin
 from databricks.labs.ucx.framework.crawlers import CrawlerBase
-from databricks.labs.ucx.framework.owners import Ownership
+from databricks.labs.ucx.framework.owners import Ownership, AdministratorLocator
 from databricks.labs.ucx.framework.utils import escape_sql_identifier
 
 logger = logging.getLogger(__name__)
@@ -191,6 +191,9 @@ class ClusterOwnership(Ownership[ClusterInfo]):
     This is the cluster creator (if known).
     """
 
+    def __init__(self, administrator_locator: AdministratorLocator):
+        super().__init__(administrator_locator, ClusterInfo)
+
     def _maybe_direct_owner(self, record: ClusterInfo) -> str | None:
         return record.creator
 
@@ -256,5 +259,8 @@ class ClusterPolicyOwnership(Ownership[PolicyInfo]):
     This is the creator of the cluster policy (if known).
     """
 
+    def __init__(self, administrator_locator: AdministratorLocator):
+        super().__init__(administrator_locator, PolicyInfo)
+
     def _maybe_direct_owner(self, record: PolicyInfo) -> str | None:
         return record.creator

src/databricks/labs/ucx/assessment/jobs.py

@@ -26,7 +26,7 @@
 from databricks.labs.ucx.assessment.clusters import CheckClusterMixin
 from databricks.labs.ucx.assessment.crawlers import spark_version_compatibility
 from databricks.labs.ucx.framework.crawlers import CrawlerBase
-from databricks.labs.ucx.framework.owners import Ownership
+from databricks.labs.ucx.framework.owners import Ownership, AdministratorLocator
 from databricks.labs.ucx.framework.utils import escape_sql_identifier
 
 logger = logging.getLogger(__name__)
@@ -158,6 +158,9 @@ class JobOwnership(Ownership[JobInfo]):
     This is the job creator (if known).
     """
 
+    def __init__(self, administrator_locator: AdministratorLocator):
+        super().__init__(administrator_locator, JobInfo)
+
     def _maybe_direct_owner(self, record: JobInfo) -> str | None:
         return record.creator
 

src/databricks/labs/ucx/assessment/pipelines.py

@@ -10,7 +10,7 @@
 
 from databricks.labs.ucx.assessment.clusters import CheckClusterMixin
 from databricks.labs.ucx.framework.crawlers import CrawlerBase
-from databricks.labs.ucx.framework.owners import Ownership
+from databricks.labs.ucx.framework.owners import Ownership, AdministratorLocator
 from databricks.labs.ucx.framework.utils import escape_sql_identifier
 
 logger = logging.getLogger(__name__)
@@ -86,5 +86,8 @@ class PipelineOwnership(Ownership[PipelineInfo]):
     This is the pipeline creator (if known).
     """
 
+    def __init__(self, administrator_locator: AdministratorLocator):
+        super().__init__(administrator_locator, PipelineInfo)
+
     def _maybe_direct_owner(self, record: PipelineInfo) -> str | None:
         return record.creator_name

src/databricks/labs/ucx/contexts/application.py

@@ -28,7 +28,12 @@
 from databricks.labs.ucx.assessment.export import AssessmentExporter
 from databricks.labs.ucx.aws.credentials import CredentialManager
 from databricks.labs.ucx.config import WorkspaceConfig
-from databricks.labs.ucx.framework.owners import AdministratorLocator, WorkspacePathOwnership, LegacyQueryOwnership
+from databricks.labs.ucx.framework.owners import (
+    AdministratorLocator,
+    WorkspacePathOwnership,
+    Ownership,
+    LegacyQueryOwnership,
+)
 from databricks.labs.ucx.hive_metastore import ExternalLocations, MountsCrawler, TablesCrawler
 from databricks.labs.ucx.hive_metastore.catalog_schema import CatalogSchema
 from databricks.labs.ucx.hive_metastore.grants import (
@@ -571,6 +576,19 @@ def migration_recon(self) -> MigrationRecon:
     def administrator_locator(self) -> AdministratorLocator:
         return AdministratorLocator(self.workspace_client)
 
+    @cached_property
+    def ownership_factory(self) -> Callable[[type], Ownership]:
+        # ensure registration of Ownerships
+        names_with_ownership = [name for name in dir(GlobalContext) if "ownership" in name]
+        for name in names_with_ownership:
+            if name == "ownership_factory":
+                continue
+            prop = getattr(GlobalContext, name)
+            if not isinstance(prop, cached_property):
+                continue
+            _ = getattr(self, name)
+        return Ownership.for_record_type
+
 
 class CliContext(GlobalContext, abc.ABC):
     @cached_property

src/databricks/labs/ucx/framework/owners.py

@@ -1,9 +1,11 @@
+from __future__ import annotations
 import logging
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterable, Sequence
+from dataclasses import dataclass
 from datetime import timedelta
 from functools import cached_property
-from typing import Generic, TypeVar, final
+from typing import Generic, TypeVar, final, cast
 
 from databricks.labs.blueprint.paths import WorkspacePath
 from databricks.sdk import WorkspaceClient
@@ -169,8 +171,15 @@ def get_workspace_administrator(self) -> str:
 class Ownership(ABC, Generic[Record]):
     """Determine an owner for a given type of object."""
 
-    def __init__(self, administrator_locator: AdministratorLocator) -> None:
+    _factories: dict[type, Ownership] = {}
+
+    @classmethod
+    def for_record_type(cls, record_type: type) -> Ownership[Record]:
+        return cast(Ownership[Record], cls._factories[record_type])
+
+    def __init__(self, administrator_locator: AdministratorLocator, record_type: type) -> None:
         self._administrator_locator = administrator_locator
+        self._factories[record_type] = self
 
     @final
     def owner_of(self, record: Record) -> str:
@@ -198,7 +207,7 @@ def _maybe_direct_owner(self, record: Record) -> str | None:
 
 class WorkspacePathOwnership(Ownership[WorkspacePath]):
     def __init__(self, administrator_locator: AdministratorLocator, ws: WorkspaceClient) -> None:
-        super().__init__(administrator_locator)
+        super().__init__(administrator_locator, WorkspacePath)
         self._ws = ws
 
     def owner_of_path(self, path: str) -> str:
@@ -242,14 +251,19 @@ def _infer_from_first_can_manage(object_permissions):
         return None
 
 
-class LegacyQueryOwnership(Ownership[str]):
+@dataclass
+class LegacyQueryPath:
+    path: str
+
+
+class LegacyQueryOwnership(Ownership[LegacyQueryPath]):
     def __init__(self, administrator_locator: AdministratorLocator, workspace_client: WorkspaceClient) -> None:
-        super().__init__(administrator_locator)
+        super().__init__(administrator_locator, LegacyQueryPath)
         self._workspace_client = workspace_client
 
-    def _maybe_direct_owner(self, record: str) -> str | None:
+    def _maybe_direct_owner(self, record: LegacyQueryPath) -> str | None:
         try:
-            legacy_query = self._workspace_client.queries.get(record)
+            legacy_query = self._workspace_client.queries.get(record.path)
             return legacy_query.owner_user_name
         except NotFound:
             return None

src/databricks/labs/ucx/hive_metastore/grants.py

@@ -32,7 +32,7 @@
     StoragePermissionMapping,
 )
 from databricks.labs.ucx.framework.crawlers import CrawlerBase
-from databricks.labs.ucx.framework.owners import Ownership
+from databricks.labs.ucx.framework.owners import Ownership, AdministratorLocator
 from databricks.labs.ucx.framework.utils import escape_sql_identifier
 from databricks.labs.ucx.hive_metastore import TablesCrawler
 from databricks.labs.ucx.hive_metastore.locations import (
@@ -404,6 +404,9 @@ class GrantOwnership(Ownership[Grant]):
     At the present we can't determine a specific owner for grants.
     """
 
+    def __init__(self, administrator_locator: AdministratorLocator):
+        super().__init__(administrator_locator, Grant)
+
     def _maybe_direct_owner(self, record: Grant) -> None:
         return None
 

src/databricks/labs/ucx/hive_metastore/ownership.py

@@ -6,6 +6,7 @@
     AdministratorLocator,
     LegacyQueryOwnership,
     WorkspacePathOwnership,
+    LegacyQueryPath,
 )
 from databricks.labs.ucx.hive_metastore import TablesCrawler
 from databricks.labs.ucx.hive_metastore.grants import GrantsCrawler
@@ -33,7 +34,7 @@ def __init__(
         legacy_query_ownership: LegacyQueryOwnership,
         workspace_path_ownership: WorkspacePathOwnership,
     ) -> None:
-        super().__init__(administrator_locator)
+        super().__init__(administrator_locator, Table)
         self._grants_crawler = grants_crawler
         self._used_tables_in_paths = used_tables_in_paths
         self._used_tables_in_queries = used_tables_in_queries
@@ -54,7 +55,7 @@ def _maybe_from_sources(self, record: Table) -> str | None:
         if not used_table.is_write:
             return None
         if used_table.source_type == 'QUERY' and used_table.query_id:
-            return self._legacy_query_ownership.owner_of(used_table.query_id)
+            return self._legacy_query_ownership.owner_of(LegacyQueryPath(used_table.query_id))
         if used_table.source_type in {'NOTEBOOK', 'FILE'}:
             return self._workspace_path_ownership.owner_of_path(used_table.source_id)
         logger.warning(f"Unknown source type {used_table.source_type} for {used_table.source_id}")
@@ -92,7 +93,7 @@ class TableMigrationOwnership(Ownership[TableMigrationStatus]):
     """
 
     def __init__(self, tables_crawler: TablesCrawler, table_ownership: TableOwnership) -> None:
-        super().__init__(table_ownership._administrator_locator)  # TODO: Fix this
+        super().__init__(table_ownership._administrator_locator, TableMigrationStatus)  # TODO: Fix this
         self._tables_crawler = tables_crawler
         self._table_ownership = table_ownership
         self._indexed_tables: dict[tuple[str, str], Table] | None = None

src/databricks/labs/ucx/hive_metastore/udfs.py

@@ -9,7 +9,7 @@
 from databricks.sdk.errors import Unknown, NotFound
 
 from databricks.labs.ucx.framework.crawlers import CrawlerBase
-from databricks.labs.ucx.framework.owners import Ownership
+from databricks.labs.ucx.framework.owners import Ownership, AdministratorLocator
 from databricks.labs.ucx.framework.utils import escape_sql_identifier
 
 logger = logging.getLogger(__name__)
@@ -151,5 +151,8 @@ class UdfOwnership(Ownership[Udf]):
     At the present we don't determine a specific owner for UDFs.
     """
 
+    def __init__(self, administrator_locator: AdministratorLocator):
+        super().__init__(administrator_locator, Udf)
+
     def _maybe_direct_owner(self, record: Udf) -> None:
         return None

src/databricks/labs/ucx/source_code/directfs_access.py

@@ -15,6 +15,7 @@
     AdministratorLocator,
     WorkspacePathOwnership,
     LegacyQueryOwnership,
+    LegacyQueryPath,
 )
 from databricks.labs.ucx.framework.utils import escape_sql_identifier
 from databricks.labs.ucx.source_code.base import DirectFsAccess
@@ -81,14 +82,14 @@ def __init__(
         legacy_query_ownership: LegacyQueryOwnership,
         workspace_client: WorkspaceClient,
     ) -> None:
-        super().__init__(administrator_locator)
+        super().__init__(administrator_locator, DirectFsAccess)
         self._workspace_path_ownership = workspace_path_ownership
         self._legacy_query_ownership = legacy_query_ownership
         self._workspace_client = workspace_client
 
     def _maybe_direct_owner(self, record: DirectFsAccess) -> str | None:
         if record.source_type == 'QUERY' and record.query_id:
-            return self._legacy_query_ownership.owner_of(record.query_id)
+            return self._legacy_query_ownership.owner_of(LegacyQueryPath(record.query_id))
         if record.source_type in {'NOTEBOOK', 'FILE'}:
             return self._notebook_owner(record)
         logger.warning(f"Unknown source type {record.source_type} for {record.source_id}")

tests/unit/contexts/test_application.py

@@ -1,12 +1,18 @@
 from unittest.mock import create_autospec
 
 import pytest
+from databricks.labs.blueprint.paths import WorkspacePath
 from databricks.labs.lsql.backends import MockBackend
 
 from databricks.labs.ucx.contexts.application import GlobalContext
 from databricks.labs.ucx.contexts.workspace_cli import LocalCheckoutContext
-from databricks.labs.ucx.hive_metastore.table_migration_status import TableMigrationIndex
+from databricks.labs.ucx.framework.owners import Ownership
+from databricks.labs.ucx.hive_metastore.grants import Grant
+from databricks.labs.ucx.hive_metastore.table_migration_status import TableMigrationIndex, TableMigrationStatus
 from databricks.labs.ucx.hive_metastore.table_migrate import TablesMigrator
+from databricks.labs.ucx.hive_metastore.tables import Table
+from databricks.labs.ucx.hive_metastore.udfs import Udf
+from databricks.labs.ucx.source_code.base import DirectFsAccess
 from databricks.labs.ucx.source_code.linters.context import LinterContext
 from tests.unit import mock_workspace_client
 
@@ -23,6 +29,7 @@
         "used_tables_crawler_for_paths",
         "used_tables_crawler_for_queries",
         "verify_has_ucx_catalog",
+        "ownership_factory",
     ],
 )
 def test_global_context_attributes_not_none(attribute: str) -> None:
@@ -44,3 +51,20 @@ def test_local_context_attributes_not_none(attribute: str) -> None:
     ctx.replace(languages=LinterContext(TableMigrationIndex([])), tables_migrator=tables_migrator)
     assert hasattr(ctx, attribute)
     assert getattr(ctx, attribute) is not None
+
+
+@pytest.mark.parametrize(
+    "record_type",
+    [
+        DirectFsAccess,
+        WorkspacePath,
+        Grant,
+        Table,
+        Udf,
+        TableMigrationStatus,
+    ],
+)
+def test_ownership_factory_succeeds(record_type: type):
+    ctx = GlobalContext().replace(workspace_client=mock_workspace_client(), sql_backend=MockBackend())
+    ownership = ctx.ownership_factory(record_type)
+    assert isinstance(ownership, Ownership)

tests/unit/framework/test_owners.py

@@ -23,7 +23,7 @@ def __init__(
         owner_fn: Callable[[Record], str | None] = lambda _: None,
     ):
         mock_admin_locator = create_autospec(AdministratorLocator)  # pylint: disable=mock-no-usage
-        super().__init__(mock_admin_locator)
+        super().__init__(mock_admin_locator, object)
         self._owner_fn = owner_fn
         self.mock_admin_locator = mock_admin_locator