Use LegacyQueryPath dataclass as record for LegacyQueryOwnership

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

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}")

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}")