Add instance pool id to WorkspaceConfig (#1087)

Author: qziyuan

src/databricks/labs/ucx/install.py

@@ -201,7 +201,9 @@ def _configure_new_installation(self) -> WorkspaceConfig:
         log_level = self._prompts.question("Log level", default="INFO").upper()
         num_threads = int(self._prompts.question("Number of threads", default="8", valid_number=True))
 
-        policy_id, instance_profile, spark_conf_dict = self._policy_installer.create(inventory_database)
+        policy_id, instance_profile, spark_conf_dict, instance_pool_id = self._policy_installer.create(
+            inventory_database
+        )
 
         # Check if terraform is being used
         is_terraform_used = self._prompts.confirm("Do you use Terraform to deploy your infrastructure?")
@@ -220,6 +222,7 @@ def _configure_new_installation(self) -> WorkspaceConfig:
             instance_profile=instance_profile,
             spark_conf=spark_conf_dict,
             policy_id=policy_id,
+            instance_pool_id=instance_pool_id,
             is_terraform_used=is_terraform_used,
             include_databases=self._select_databases(),
         )

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

@@ -22,7 +22,7 @@ def __init__(self, installation: Installation, ws: WorkspaceClient, prompts: Pro
     def _policy_config(value: str):
         return {"type": "fixed", "value": value}
 
-    def create(self, inventory_database: str) -> tuple[str, str, dict]:
+    def create(self, inventory_database: str) -> tuple[str, str, dict, str | None]:
         instance_profile = ""
         spark_conf_dict = {}
         # get instance pool id to be put into the cluster policy
@@ -52,7 +52,7 @@ def create(self, inventory_database: str) -> tuple[str, str, dict]:
                 logger.info(f"Cluster policy {policy_name} already present, reusing the same.")
                 policy_id = policy.policy_id
                 assert policy_id is not None
-                return policy_id, instance_profile, spark_conf_dict
+                return policy_id, instance_profile, spark_conf_dict, instance_pool_id
         logger.info("Creating UCX cluster policy.")
         policy_id = self._ws.cluster_policies.create(
             name=policy_name,
@@ -64,6 +64,7 @@ def create(self, inventory_database: str) -> tuple[str, str, dict]:
             policy_id,
             instance_profile,
             spark_conf_dict,
+            instance_pool_id,
         )
 
     def _get_instance_pool_id(self) -> str | None:
@@ -72,7 +73,7 @@ def _get_instance_pool_id(self) -> str | None:
                 "Instance pool id to be set in cluster policy for all workflow clusters", default="None"
             )
         except OSError:
-            # when unit test v0.15.0_added_cluster_policy.py MockPromots cannot be injected to ClusterPolicyInstaller
+            # when unit test v0.15.0_added_cluster_policy.py MockPrompts cannot be injected to ClusterPolicyInstaller
             # return None to pass the test
             return None
         if instance_pool_id.lower() == "none":

src/databricks/labs/ucx/upgrades/v0.15.0_added_cluster_policy.py

@@ -16,7 +16,7 @@
 def upgrade(installation: Installation, ws: WorkspaceClient):
     config = installation.load(WorkspaceConfig)
     policy_installer = ClusterPolicyInstaller(installation, ws, Prompts())
-    config.policy_id, _, _ = policy_installer.create(config.inventory_database)
+    config.policy_id, _, _, _ = policy_installer.create(config.inventory_database)
     installation.save(config)
     states = InstallState.from_installation(installation)
     assert config.policy_id is not None

tests/integration/test_installation.py

@@ -2,6 +2,7 @@
 import json
 import logging
 import os.path
+import sys
 from collections.abc import Callable
 from dataclasses import replace
 from datetime import timedelta
@@ -474,9 +475,9 @@ def test_check_inventory_database_exists(ws, new_installation):
 def test_table_migration_job(  # pylint: disable=too-many-locals
     ws, new_installation, make_catalog, make_schema, make_table, env_or_skip, make_random, make_dbfs_data_copy
 ):
-    # skip this test if not in nightly test job: TEST_NIGHTLY is missing or is not set to "true"
-    if env_or_skip("TEST_NIGHTLY").lower() != "true":
-        pytest.skip("TEST_NIGHTLY is not true")
+    # skip this test if not in nightly test job or debug mode
+    if os.path.basename(sys.argv[0]) not in {"_jb_pytest_runner.py", "testlauncher.py"}:
+        env_or_skip("TEST_NIGHTLY")
     # create external and managed tables to be migrated
     src_schema = make_schema(catalog_name="hive_metastore")
     src_managed_table = make_table(schema_name=src_schema.name)

tests/unit/installer/test_policy.py

@@ -44,7 +44,7 @@ def test_cluster_policy_definition_present_reuse():
     ws, prompts = common()
 
     policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
-    policy_id, _, _ = policy_installer.create('ucx')
+    policy_id, _, _, _ = policy_installer.create('ucx')
     assert policy_id is not None
     ws.cluster_policies.create.assert_not_called()
 
@@ -73,7 +73,7 @@ def test_cluster_policy_definition_azure_hms():
         )
     ]
     policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
-    policy_id, _, _ = policy_installer.create('ucx')
+    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"},
@@ -114,7 +114,7 @@ def test_cluster_policy_definition_aws_glue():
     ]
 
     policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
-    policy_id, instance_profile, _ = policy_installer.create('ucx')
+    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"},
@@ -155,7 +155,7 @@ def test_cluster_policy_definition_gcp():
     ]
 
     policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
-    policy_id, instance_profile, _ = policy_installer.create('ucx')
+    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"},
@@ -257,7 +257,7 @@ def test_cluster_policy_definition_azure_hms_warehouse():
         }
     )
     policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
-    policy_id, _, _ = policy_installer.create('ucx')
+    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"},
@@ -309,7 +309,7 @@ def test_cluster_policy_definition_aws_glue_warehouse():
         }
     )
     policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
-    policy_id, instance_profile, _ = policy_installer.create('ucx')
+    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"},
@@ -364,7 +364,7 @@ def test_cluster_policy_definition_gcp_hms_warehouse():
         }
     )
     policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
-    policy_id, _, _ = policy_installer.create('ucx')
+    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"},
@@ -405,7 +405,7 @@ def test_cluster_policy_definition_empty_config():
     ]
 
     policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
-    policy_id, _, _ = policy_installer.create('ucx')
+    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"},
@@ -431,7 +431,9 @@ def test_cluster_policy_instance_pool():
     ws.config.is_gcp = False
 
     policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
-    policy_installer.create('ucx')
+    _, _, _, instance_pool_id = policy_installer.create('ucx')
+
+    assert instance_pool_id == "instance_pool_1"
 
     policy_expected = {
         "spark_version": {"type": "fixed", "value": "14.2.x-scala2.12"},
@@ -452,7 +454,8 @@ def test_cluster_policy_instance_pool():
         "node_type_id": {"type": "fixed", "value": "Standard_F4s"},
         "aws_attributes.availability": {"type": "fixed", "value": "ON_DEMAND"},
     }
-    policy_installer.create('ucx')
+    _, _, _, instance_pool_id = policy_installer.create('ucx')
+    assert instance_pool_id is None
     ws.cluster_policies.create.assert_called_with(
         name="Unity Catalog Migration (ucx) ([email protected])",
         definition=json.dumps(policy_expected),