Add a check for external metastore in SQL warehouse configuration (#1046)

Author: nkvuong

src/databricks/labs/ucx/installer/policy.py

@@ -7,6 +7,7 @@
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import NotFound
 from databricks.sdk.service import compute
+from databricks.sdk.service.sql import GetWorkspaceWarehouseConfigResponse
 
 logger = logging.getLogger(__name__)
 
@@ -34,6 +35,14 @@ def create(self, inventory_database: str) -> tuple[str, str, dict]:
             if len(cluster_policies) >= 1:
                 cluster_policy = json.loads(self._prompts.choice_from_dict("Choose a cluster policy", cluster_policies))
                 instance_profile, spark_conf_dict = self._extract_external_hive_metastore_conf(cluster_policy)
+        else:
+            warehouse_config = self._get_warehouse_config_with_external_hive_metastore()
+            if warehouse_config and self._prompts.confirm(
+                "We have identified the workspace warehouse is set up for an external metastore"
+                "Would you like to set UCX to connect to the external metastore?"
+            ):
+                logger.info("Setting up an external metastore")
+                instance_profile, spark_conf_dict = self._extract_external_hive_metastore_sql_conf(warehouse_config)
         policy_name = f"Unity Catalog Migration ({inventory_database}) ({self._ws.current_user.me().user_name})"
         policies = self._ws.cluster_policies.list()
         for policy in policies:
@@ -63,12 +72,12 @@ def _definition(self, conf: dict, instance_profile: str | None) -> str:
         for key, value in conf.items():
             policy_definition[f"spark_conf.{key}"] = self._policy_config(value)
         if self._ws.config.is_aws:
+            if instance_profile:
+                policy_definition["aws_attributes.instance_profile_arn"] = self._policy_config(instance_profile)
             policy_definition["aws_attributes.availability"] = self._policy_config(
                 compute.AwsAvailability.ON_DEMAND.value
             )
-            if instance_profile:
-                policy_definition["aws_attributes.instance_profile_arn"] = self._policy_config(instance_profile)
-        elif self._ws.config.is_azure:  # pylint: disable=confusing-consecutive-elif
+        elif self._ws.config.is_azure:
             policy_definition["azure_attributes.availability"] = self._policy_config(
                 compute.AzureAvailability.ON_DEMAND_AZURE.value
             )
@@ -107,6 +116,41 @@ def _get_cluster_policies_with_external_hive_metastores(self):
                     yield policy
                     break
 
+    @staticmethod
+    def _extract_external_hive_metastore_sql_conf(sql_config: GetWorkspaceWarehouseConfigResponse):
+        spark_conf_dict: dict[str, str] = {}
+        instance_profile = None
+        if sql_config.instance_profile_arn is not None:
+            instance_profile = sql_config.instance_profile_arn
+            logger.info(f"Instance Profile is Set to {instance_profile}")
+        if sql_config.data_access_config is None:
+            return instance_profile, spark_conf_dict
+        for conf in sql_config.data_access_config:
+            if conf.key is None:
+                continue
+            if conf.value is None:
+                continue
+            if (
+                conf.key.startswith("spark.sql.hive.metastore")
+                or conf.key.startswith("spark.hadoop.javax.jdo.option")
+                or conf.key.startswith("spark.databricks.hive.metastore")
+                or conf.key.startswith("spark.hadoop.hive.metastore.glue")
+            ):
+                spark_conf_dict[conf.key] = conf.value
+        return instance_profile, spark_conf_dict
+
+    def _get_warehouse_config_with_external_hive_metastore(self) -> GetWorkspaceWarehouseConfigResponse | None:
+        sql_config = self._ws.warehouses.get_workspace_warehouse_config()
+        if sql_config.data_access_config is None:
+            return None
+        for conf in sql_config.data_access_config:
+            if conf.key is None:
+                continue
+            is_glue = conf.key.startswith("spark.databricks.hive.metastore.glueCatalog.enabled")
+            if conf.key.startswith("spark.sql.hive.metastore") or is_glue:
+                return sql_config
+        return None
+
     def update_job_policy(self, state: InstallState, policy_id: str):
         if not state.jobs:
             logger.error("No jobs found in states")
@@ -118,7 +162,8 @@ def update_job_policy(self, state: InstallState, policy_id: str):
                 assert job.job_id is not None
                 assert job_settings is not None
                 if job_settings.job_clusters is None:
-                    # if job_clusters is None, it means override cluster is being set and hence policy should not be applied
+                    # if job_clusters is None, it means override cluster is being set
+                    # and hence policy should not be applied
                     return
             except NotFound:
                 logger.error(f"Job id {job_id} not found. Please check if the job is present in the workspace")

tests/unit/installer/test_policy.py

@@ -9,6 +9,10 @@
 from databricks.sdk.service import iam
 from databricks.sdk.service.compute import ClusterSpec, Policy
 from databricks.sdk.service.jobs import Job, JobCluster, JobSettings
+from databricks.sdk.service.sql import (
+    EndpointConfPair,
+    GetWorkspaceWarehouseConfigResponse,
+)
 
 from databricks.labs.ucx.installer.policy import ClusterPolicyInstaller
 
@@ -67,12 +71,6 @@ def test_cluster_policy_definition_azure_hms():
             description="Custom cluster policy for Unity Catalog Migration (UCX)",
         )
     ]
-    prompts = MockPrompts(
-        {
-            r".*We have identified one or more cluster.*": "Yes",
-            r".*Choose a cluster policy.*": "0",
-        }
-    )
     policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
     policy_id, _, _ = policy_installer.create('ucx')
     policy_definition_actual = {
@@ -120,8 +118,8 @@ def test_cluster_policy_definition_aws_glue():
         "spark_version": {"type": "fixed", "value": "14.2.x-scala2.12"},
         "node_type_id": {"type": "fixed", "value": "Standard_F4s"},
         "spark_conf.spark.databricks.hive.metastore.glueCatalog.enabled": {"type": "fixed", "value": "true"},
-        "aws_attributes.availability": {"type": "fixed", "value": "ON_DEMAND"},
         "aws_attributes.instance_profile_arn": {"type": "fixed", "value": "role_arn_1"},
+        "aws_attributes.availability": {"type": "fixed", "value": "ON_DEMAND"},
     }
     assert policy_id == "foo1"
     assert instance_profile == "role_arn_1"
@@ -221,3 +219,198 @@ def test_update_job_policy():
     ws.jobs.get.return_value = job
     policy_installer.update_job_policy(states, 'foobar')
     ws.jobs.update.assert_called_with(123, new_settings=job_setting)
+
+
+def test_cluster_policy_definition_azure_hms_warehouse():
+    ws, _ = common()
+    ws.config.is_aws = False
+    ws.config.is_azure = True
+    ws.config.is_gcp = False
+    endpoint_conf = [
+        EndpointConfPair("spark.hadoop.javax.jdo.option.ConnectionURL", "url"),
+        EndpointConfPair("spark.hadoop.javax.jdo.option.ConnectionUserName", "user1"),
+        EndpointConfPair("spark.hadoop.javax.jdo.option.ConnectionPassword", "pwd"),
+        EndpointConfPair("spark.hadoop.javax.jdo.option.ConnectionDriverName", "SQLServerDriver"),
+        EndpointConfPair("spark.sql.hive.metastore.version", "0.13"),
+        EndpointConfPair("spark.sql.hive.metastore.jars", "jar1"),
+    ]
+
+    ws.warehouses.get_workspace_warehouse_config.return_value = GetWorkspaceWarehouseConfigResponse(
+        data_access_config=endpoint_conf
+    )
+
+    ws.cluster_policies.list.return_value = [
+        Policy(
+            policy_id="id1",
+            name="foo",
+            definition=json.dumps({}),
+            description="Custom cluster policy for Unity Catalog Migration (UCX)",
+        )
+    ]
+
+    prompts = MockPrompts(
+        {
+            r".*We have identified one or more cluster.*": "No",
+            r".*We have identified the workspace warehouse.*": "Yes",
+        }
+    )
+    policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
+    policy_id, _, _ = policy_installer.create('ucx')
+    policy_definition_actual = {
+        "spark_version": {"type": "fixed", "value": "14.2.x-scala2.12"},
+        "node_type_id": {"type": "fixed", "value": "Standard_F4s"},
+        "spark_conf.spark.hadoop.javax.jdo.option.ConnectionURL": {"type": "fixed", "value": "url"},
+        "spark_conf.spark.hadoop.javax.jdo.option.ConnectionUserName": {"type": "fixed", "value": "user1"},
+        "spark_conf.spark.hadoop.javax.jdo.option.ConnectionPassword": {"type": "fixed", "value": "pwd"},
+        "spark_conf.spark.hadoop.javax.jdo.option.ConnectionDriverName": {"type": "fixed", "value": "SQLServerDriver"},
+        "spark_conf.spark.sql.hive.metastore.version": {"type": "fixed", "value": "0.13"},
+        "spark_conf.spark.sql.hive.metastore.jars": {"type": "fixed", "value": "jar1"},
+        "azure_attributes.availability": {"type": "fixed", "value": "ON_DEMAND_AZURE"},
+    }
+    assert policy_id == "foo1"
+
+    ws.cluster_policies.create.assert_called_with(
+        name="Unity Catalog Migration (ucx) ([email protected])",
+        definition=json.dumps(policy_definition_actual),
+        description="Custom cluster policy for Unity Catalog Migration (UCX)",
+    )
+
+
+def test_cluster_policy_definition_aws_glue_warehouse():
+    ws, _ = common()
+    ws.config.is_aws = True
+    ws.config.is_azure = False
+    ws.config.is_gcp = False
+    endpoint_conf = [
+        EndpointConfPair("spark.databricks.hive.metastore.glueCatalog.enabled", "true"),
+    ]
+
+    ws.warehouses.get_workspace_warehouse_config.return_value = GetWorkspaceWarehouseConfigResponse(
+        data_access_config=endpoint_conf,
+        instance_profile_arn="role_arn_1",
+    )
+
+    ws.cluster_policies.list.return_value = [
+        Policy(
+            policy_id="id1",
+            name="foo",
+            definition=json.dumps({}),
+            description="Custom cluster policy for Unity Catalog Migration (UCX)",
+        )
+    ]
+
+    prompts = MockPrompts(
+        {
+            r".*We have identified one or more cluster.*": "No",
+            r".*We have identified the workspace warehouse.*": "Yes",
+        }
+    )
+    policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
+    policy_id, instance_profile, _ = policy_installer.create('ucx')
+    policy_definition_actual = {
+        "spark_version": {"type": "fixed", "value": "14.2.x-scala2.12"},
+        "node_type_id": {"type": "fixed", "value": "Standard_F4s"},
+        "spark_conf.spark.databricks.hive.metastore.glueCatalog.enabled": {"type": "fixed", "value": "true"},
+        "aws_attributes.instance_profile_arn": {"type": "fixed", "value": "role_arn_1"},
+        "aws_attributes.availability": {"type": "fixed", "value": "ON_DEMAND"},
+    }
+    assert policy_id == "foo1"
+    assert instance_profile == "role_arn_1"
+
+    ws.cluster_policies.create.assert_called_with(
+        name="Unity Catalog Migration (ucx) ([email protected])",
+        definition=json.dumps(policy_definition_actual),
+        description="Custom cluster policy for Unity Catalog Migration (UCX)",
+    )
+
+
+def test_cluster_policy_definition_gcp_hms_warehouse():
+    ws, _ = common()
+    ws.config.is_aws = False
+    ws.config.is_azure = False
+    ws.config.is_gcp = True
+    endpoint_conf = [
+        EndpointConfPair(None, None),
+        EndpointConfPair("random", None),
+        EndpointConfPair("spark.hadoop.javax.jdo.option.ConnectionURL", "url"),
+        EndpointConfPair("spark.hadoop.javax.jdo.option.ConnectionUserName", "user1"),
+        EndpointConfPair("spark.hadoop.javax.jdo.option.ConnectionPassword", "pwd"),
+        EndpointConfPair("spark.hadoop.javax.jdo.option.ConnectionDriverName", "SQLServerDriver"),
+        EndpointConfPair("spark.sql.hive.metastore.version", "0.13"),
+        EndpointConfPair("spark.sql.hive.metastore.jars", "jar1"),
+    ]
+
+    ws.warehouses.get_workspace_warehouse_config.return_value = GetWorkspaceWarehouseConfigResponse(
+        data_access_config=endpoint_conf
+    )
+
+    ws.cluster_policies.list.return_value = [
+        Policy(
+            policy_id="id1",
+            name="foo",
+            definition=json.dumps({}),
+            description="Custom cluster policy for Unity Catalog Migration (UCX)",
+        )
+    ]
+
+    prompts = MockPrompts(
+        {
+            r".*We have identified one or more cluster.*": "No",
+            r".*We have identified the workspace warehouse.*": "Yes",
+        }
+    )
+    policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
+    policy_id, _, _ = policy_installer.create('ucx')
+    policy_definition_actual = {
+        "spark_version": {"type": "fixed", "value": "14.2.x-scala2.12"},
+        "node_type_id": {"type": "fixed", "value": "Standard_F4s"},
+        "spark_conf.spark.hadoop.javax.jdo.option.ConnectionURL": {"type": "fixed", "value": "url"},
+        "spark_conf.spark.hadoop.javax.jdo.option.ConnectionUserName": {"type": "fixed", "value": "user1"},
+        "spark_conf.spark.hadoop.javax.jdo.option.ConnectionPassword": {"type": "fixed", "value": "pwd"},
+        "spark_conf.spark.hadoop.javax.jdo.option.ConnectionDriverName": {"type": "fixed", "value": "SQLServerDriver"},
+        "spark_conf.spark.sql.hive.metastore.version": {"type": "fixed", "value": "0.13"},
+        "spark_conf.spark.sql.hive.metastore.jars": {"type": "fixed", "value": "jar1"},
+        "gcp_attributes.availability": {"type": "fixed", "value": "ON_DEMAND_GCP"},
+    }
+    assert policy_id == "foo1"
+
+    ws.cluster_policies.create.assert_called_with(
+        name="Unity Catalog Migration (ucx) ([email protected])",
+        definition=json.dumps(policy_definition_actual),
+        description="Custom cluster policy for Unity Catalog Migration (UCX)",
+    )
+
+
+def test_cluster_policy_definition_empty_config():
+    ws, prompts = common()
+    ws.config.is_aws = True
+    ws.config.is_azure = False
+    ws.config.is_gcp = False
+
+    ws.warehouses.get_workspace_warehouse_config.return_value = GetWorkspaceWarehouseConfigResponse(
+        data_access_config=None
+    )
+
+    ws.cluster_policies.list.return_value = [
+        Policy(
+            policy_id="id1",
+            name="foo",
+            definition=json.dumps({}),
+            description="Custom cluster policy for Unity Catalog Migration (UCX)",
+        )
+    ]
+
+    policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
+    policy_id, _, _ = policy_installer.create('ucx')
+    policy_definition_actual = {
+        "spark_version": {"type": "fixed", "value": "14.2.x-scala2.12"},
+        "node_type_id": {"type": "fixed", "value": "Standard_F4s"},
+        "aws_attributes.availability": {"type": "fixed", "value": "ON_DEMAND"},
+    }
+    assert policy_id == "foo1"
+
+    ws.cluster_policies.create.assert_called_with(
+        name="Unity Catalog Migration (ucx) ([email protected])",
+        definition=json.dumps(policy_definition_actual),
+        description="Custom cluster policy for Unity Catalog Migration (UCX)",
+    )