Verify metastore exists in current workspace (#735)

Author: pritishpai

labs.yml

@@ -65,3 +65,4 @@ commands:
         description: Table to revert (if left blank all tables in the schema will be reverted). Requires schema parameter to be specified.)
       - name: delete_managed
         description: Revert and delete managed tables
+

src/databricks/labs/ucx/runtime.py

@@ -22,6 +22,7 @@
 from databricks.labs.ucx.workspace_access.generic import WorkspaceListing
 from databricks.labs.ucx.workspace_access.groups import GroupManager
 from databricks.labs.ucx.workspace_access.manager import PermissionManager
+from databricks.labs.ucx.workspace_access.verification import VerifyHasMetastore
 
 logger = logging.getLogger(__name__)
 
@@ -250,6 +251,10 @@ def rename_workspace_local_groups(cfg: WorkspaceConfig):
     """Renames workspace local groups by adding `ucx-renamed-` prefix."""
     sql_backend = RuntimeBackend()
     ws = WorkspaceClient(config=cfg.to_databricks_config())
+    verify_has_metastore = VerifyHasMetastore(ws)
+    if verify_has_metastore.verify_metastore():
+        logger.info("Metastore exists in the workspace")
+
     group_manager = GroupManager(
         sql_backend,
         ws,

src/databricks/labs/ucx/workspace_access/verification.py

@@ -1,6 +1,8 @@
+import logging
 from typing import Literal
 
 from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors import PermissionDenied
 
 from databricks.labs.ucx.workspace_access.groups import MigrationState
 from databricks.labs.ucx.workspace_access.secrets import SecretScopesSupport
@@ -79,3 +81,36 @@ def verify_roles_and_entitlements(self, migration_state: MigrationState, target:
 
             assert base_group_info.roles == target_group_info.roles
             assert base_group_info.entitlements == target_group_info.entitlements
+
+
+class VerifyHasMetastore:
+    def __init__(self, ws: WorkspaceClient):
+        self.metastore_id: str | None = None
+        self.default_catalog_name: str | None = None
+        self.workspace_id: int | None = None
+        self._ws = ws
+
+    def verify_metastore(self):
+        """
+        Verifies if a metastore exists for a metastore
+        :param :
+        :return:
+        """
+
+        try:
+            current_metastore = self._ws.metastores.current()
+            if current_metastore:
+                self.default_catalog_name = current_metastore.default_catalog_name
+                self.metastore_id = current_metastore.metastore_id
+                self.workspace_id = current_metastore.workspace_id
+                return True
+            else:
+                raise MetastoreNotFoundError
+        except PermissionDenied:
+            logger.error("Permission Denied while trying to access metastore")
+
+
+class MetastoreNotFoundError(Exception):
+    def __init__(self, message="Metastore not found in the workspace"):
+        self.message = message
+        super().__init__(self.message)

tests/unit/workspace_access/test_verification.py

@@ -0,0 +1,50 @@
+import pytest
+from databricks.sdk.errors import PermissionDenied
+from databricks.sdk.service.catalog import MetastoreAssignment
+
+from databricks.labs.ucx.workspace_access.verification import (
+    MetastoreNotFoundError,
+    VerifyHasMetastore,
+)
+
+
+def test_validate_metastore_exists(mocker):
+    ws = mocker.patch("databricks.sdk.WorkspaceClient.__init__")
+    ws.metastores = mocker.patch("databricks.sdk.WorkspaceClient.metastores")
+    ws.metastores.current = lambda: MetastoreAssignment(
+        metastore_id="21fwef-b2345-sdas-2343-sddsvv332", workspace_id=1234567890, default_catalog_name="hive_metastore"
+    )
+    verify_metastore_obj = VerifyHasMetastore(ws)
+
+    assert verify_metastore_obj.verify_metastore() is True
+
+    assert verify_metastore_obj.metastore_id == "21fwef-b2345-sdas-2343-sddsvv332"
+    assert verify_metastore_obj.default_catalog_name == "hive_metastore"
+    assert verify_metastore_obj.workspace_id == 1234567890
+
+
+def test_validate_no_metastore_exists(mocker):
+    ws = mocker.patch("databricks.sdk.WorkspaceClient.__init__")
+    ws.metastores = mocker.patch("databricks.sdk.WorkspaceClient.metastores")
+    ws.metastores.current = mocker.patch(
+        "databricks.sdk.service.catalog.MetastoreAssignment.__init__", return_value=None
+    )
+    ws.metastores.current.return_value = None
+    ws.return_value = None
+
+    verify_metastore_obj = VerifyHasMetastore(ws)
+
+    with pytest.raises(MetastoreNotFoundError, match="Metastore not found in the workspace"):
+        verify_metastore_obj.verify_metastore()
+
+
+def test_permission_denied_error(mocker):
+    ws = mocker.patch("databricks.sdk.WorkspaceClient.__init__")
+    ws.metastores = mocker.patch("databricks.sdk.WorkspaceClient.metastores")
+    ws.metastores.current.side_effect = PermissionDenied()
+    ws.metastores.current.return_value = None
+    ws.return_value = None
+
+    verify_metastore_obj = VerifyHasMetastore(ws)
+
+    assert verify_metastore_obj.verify_metastore() is None