diff --git a/src/databricks/labs/ucx/framework/owners.py b/src/databricks/labs/ucx/framework/owners.py index c9d16fca1c..71698ab1de 100644 --- a/src/databricks/labs/ucx/framework/owners.py +++ b/src/databricks/labs/ucx/framework/owners.py @@ -2,6 +2,7 @@ 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, cast @@ -250,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, str) + 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 diff --git a/src/databricks/labs/ucx/hive_metastore/ownership.py b/src/databricks/labs/ucx/hive_metastore/ownership.py index 2a2530ef68..e1e76a629d 100644 --- a/src/databricks/labs/ucx/hive_metastore/ownership.py +++ b/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 @@ -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}") diff --git a/src/databricks/labs/ucx/source_code/directfs_access.py b/src/databricks/labs/ucx/source_code/directfs_access.py index 2e794d4b26..6e450021fc 100644 --- a/src/databricks/labs/ucx/source_code/directfs_access.py +++ b/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 @@ -88,7 +89,7 @@ def __init__( 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}")