Decouple group configuration from install.py (#714)

This PR reduces the size of `install.py` file

Author: nfx

pyproject.toml

@@ -61,8 +61,8 @@ python="3.10"
 path = ".venv"
 
 [tool.hatch.envs.default.scripts]
-test        = "pytest -n auto --cov src --cov-report=xml --timeout 10 tests/unit"
-coverage    = "pytest -n auto --cov src tests/unit --timeout 10 --cov-report=html"
+test        = "pytest -n auto --cov src --cov-report=xml --timeout 30 tests/unit"
+coverage    = "pytest -n auto --cov src tests/unit --timeout 30 --cov-report=html"
 integration = "pytest -n 10 --cov src tests/integration"
 fmt         = ["isort .",
                "black .",

src/databricks/labs/ucx/install.py

@@ -1,8 +1,6 @@
-import functools
 import json
 import logging
 import os
-import re
 import sys
 import time
 import webbrowser
@@ -47,7 +45,7 @@
 from databricks.labs.ucx.runtime import main
 from databricks.labs.ucx.workspace_access.base import Permissions
 from databricks.labs.ucx.workspace_access.generic import WorkspaceObjectInfo
-from databricks.labs.ucx.workspace_access.groups import MigratedGroup
+from databricks.labs.ucx.workspace_access.groups import ConfigureGroups, MigratedGroup
 
 TAG_STEP = "step"
 TAG_APP = "App"
@@ -380,77 +378,10 @@ def warehouse_type(_):
             )
             warehouse_id = new_warehouse.id
 
-        # Setting up group migration parameters
-        groups_config_args = {}
-
-        ask_for_group = functools.partial(self._prompts.question, validate=self._is_valid_group_str)
-        ask_for_regex = functools.partial(self._prompts.question, validate=self._validate_regex)
-        while self._prompts.confirm(
-            "Do you need to convert the workspace groups to match the account groups' name?"
-            " If the workspace groups' names match the account groups' names select no"
-            " or hit <Enter/Return>."
-        ):
-            logger.info("Setting up group name translation")
-            groups_config_args["convert_group_names"] = "yes"
-            choices = {
-                "Apply a Prefix": "prefix",
-                "Apply a Suffix": "suffix",
-                "Regex Substitution": "sub",
-                "Regex Matching": "match",
-                "Match By External ID": "external",
-                "Cancel": "cancel",
-            }
-            choice = self._prompts.choice_from_dict("Choose how to map the workspace groups:", choices, sort=False)
-            match choice:
-                case "cancel":
-                    continue
-                case "prefix":
-                    prefix = ask_for_group("Enter a prefix to add to the workspace group name")
-                    if not prefix:
-                        continue
-                    groups_config_args["workspace_group_match_regex"] = "^"
-                    groups_config_args["workspace_group_replace"] = prefix
-                case "suffix":
-                    suffix = ask_for_group("Enter a suffix to add to the workspace group name")
-                    if not suffix:
-                        continue
-                    groups_config_args["workspace_group_match_regex"] = "$"
-                    groups_config_args["workspace_group_replace"] = suffix
-                case "sub":
-                    match_value = ask_for_regex("Enter a regular expression for substitution")
-                    if not match_value:
-                        continue
-                    sub_value = ask_for_group("Enter the substitution value")
-                    if not sub_value:
-                        continue
-                    groups_config_args["workspace_group_match_regex"] = match_value
-                    groups_config_args["workspace_group_replace"] = sub_value
-                case "matching":
-                    ws_match_value = ask_for_regex("Enter a regular expression to match on the workspace group")
-                    if not ws_match_value:
-                        continue
-                    acct_match_value = ask_for_regex("Enter a regular expression to match on the account group")
-                    if not acct_match_value:
-                        continue
-                    groups_config_args["workspace_group_match_regex"] = ws_match_value
-                    groups_config_args["account_group_match_regex"] = acct_match_value
-                case "external":
-                    groups_config_args["group_match_by_external_id"] = True
-            break
-
-        selected_groups = self._prompts.question(
-            "Comma-separated list of workspace group names to migrate. If not specified, we'll use all "
-            "account-level groups with matching names to workspace-level groups",
-            default="<ALL>",
-        )
-        backup_group_prefix = ask_for_group("Backup prefix", default="db-temp-")
+        configure_groups = ConfigureGroups(self._prompts)
+        configure_groups.run()
         log_level = self._prompts.question("Log level", default="INFO").upper()
         num_threads = int(self._prompts.question("Number of threads", default="8", valid_number=True))
-        groups_config_args["backup_group_prefix"] = backup_group_prefix
-        if selected_groups != "<ALL>":
-            groups_config_args["selected"] = [x.strip() for x in selected_groups.split(",")]
-        else:
-            groups_config_args["selected"] = None
 
         # Checking for external HMS
         instance_profile = None
@@ -474,12 +405,12 @@ def warehouse_type(_):
 
         self._config = WorkspaceConfig(
             inventory_database=inventory_database,
-            workspace_group_regex=groups_config_args.get("workspace_group_regex"),
-            workspace_group_replace=groups_config_args.get("workspace_group_replace"),
-            account_group_regex=groups_config_args.get("account_group_regex"),
-            group_match_by_external_id=groups_config_args.get("group_match_by_external_id"),
-            include_group_names=groups_config_args["selected"],
-            renamed_group_prefix=groups_config_args["backup_group_prefix"],
+            workspace_group_regex=configure_groups.workspace_group_regex,
+            workspace_group_replace=configure_groups.workspace_group_replace,
+            account_group_regex=configure_groups.account_group_regex,
+            group_match_by_external_id=configure_groups.group_match_by_external_id,
+            include_group_names=configure_groups.include_group_names,
+            renamed_group_prefix=configure_groups.renamed_group_prefix,
             warehouse_id=warehouse_id,
             log_level=log_level,
             num_threads=num_threads,
@@ -643,15 +574,6 @@ def _create_debug(self, remote_wheel: str):
     def notebook_link(self, path: str) -> str:
         return f"{self._ws.config.host}/#workspace{path}"
 
-    @staticmethod
-    def _validate_regex(regex_input: str) -> bool:
-        try:
-            re.compile(regex_input)
-            return True
-        except re.error:
-            logger.error(f"{regex_input} is an invalid regular expression")
-            return False
-
     def _job_settings(self, step_name: str, remote_wheel: str):
         email_notifications = None
         if not self._current_config.override_clusters and "@" in self._my_username:
@@ -858,10 +780,6 @@ def _get_ext_hms_conf_from_policy(cluster_policy):
                 spark_conf_dict[key[11:]] = cluster_policy[key]["value"]
         return instance_profile, spark_conf_dict
 
-    @staticmethod
-    def _is_valid_group_str(group_str: str):
-        return group_str and not re.search(r"[\s#,+ \\<>;]", group_str)
-
     def latest_job_status(self) -> list[dict]:
         latest_status = []
         for step, job_id in self._state.jobs.items():

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

@@ -22,6 +22,7 @@
 
 from databricks.labs.ucx.framework.crawlers import CrawlerBase, SqlBackend
 from databricks.labs.ucx.framework.parallel import ManyError, Threads
+from databricks.labs.ucx.framework.tui import Prompts
 from databricks.labs.ucx.mixins.hardening import rate_limited
 
 logger = logging.getLogger(__name__)
@@ -73,13 +74,6 @@ def is_in_scope(self, name: str) -> bool:
         else:
             return name in self._name_to_group
 
-    def get_target_id(self, group_id: str) -> str | None:
-        group = self._id_to_group.get(group_id)
-        if group:
-            return group.external_id
-        else:
-            return None
-
     def __len__(self):
         return len(self._name_to_group)
 
@@ -539,3 +533,97 @@ def _get_strategy(
             renamed_groups_prefix=self._renamed_group_prefix,
             include_group_names=self._include_group_names,
         )
+
+
+class ConfigureGroups:
+    renamed_group_prefix = "db-temp-"
+    workspace_group_regex = None
+    workspace_group_replace = None
+    account_group_regex = None
+    group_match_by_external_id = None
+    include_group_names = None
+
+    def __init__(self, prompts: Prompts):
+        self._prompts = prompts
+        self._ask_for_group = functools.partial(self._prompts.question, validate=self._is_valid_group_str)
+        self._ask_for_regex = functools.partial(self._prompts.question, validate=self._validate_regex)
+
+    def run(self):
+        self.renamed_group_prefix = self._ask_for_group("Backup prefix", default=self.renamed_group_prefix)
+        strategy = self._prompts.choice_from_dict(
+            "Choose how to map the workspace groups:",
+            {
+                "Apply a Prefix": self._configure_prefix,
+                "Apply a Suffix": self._configure_suffix,
+                "Comma-separated list of workspace group names to migrate": self._configure_names,
+                "Match By External ID": self._configure_external,
+                "Regex Substitution": self._configure_substitution,
+                "Regex Matching": self._configure_matching,
+            },
+        )
+        strategy()
+
+    def _configure_prefix(self):
+        prefix = self._ask_for_group("Enter a prefix to add to the workspace group name")
+        if not prefix:
+            return False
+        self.workspace_group_regex = "^"
+        self.workspace_group_replace = prefix
+        return True
+
+    def _configure_suffix(self):
+        suffix = self._ask_for_group("Enter a suffix to add to the workspace group name")
+        if not suffix:
+            return False
+        self.workspace_group_regex = "$"
+        self.workspace_group_replace = suffix
+        return True
+
+    def _configure_substitution(self):
+        match_value = self._ask_for_regex("Enter a regular expression for substitution")
+        if not match_value:
+            return False
+        sub_value = self._ask_for_group("Enter the substitution value")
+        if not sub_value:
+            return False
+        self.workspace_group_regex = match_value
+        self.workspace_group_replace = sub_value
+        return True
+
+    def _configure_matching(self):
+        ws_match_value = self._ask_for_regex("Enter a regular expression to match on the workspace group")
+        if not ws_match_value:
+            return False
+        acct_match_value = self._ask_for_regex("Enter a regular expression to match on the account group")
+        if not acct_match_value:
+            return False
+        self.workspace_group_regex = ws_match_value
+        self.account_group_regex = acct_match_value
+        return True
+
+    def _configure_names(self):
+        selected_groups = self._prompts.question(
+            "Comma-separated list of workspace group names to migrate. If not specified, we'll use all "
+            "account-level groups with matching names to workspace-level groups",
+            default="<ALL>",
+        )
+        if selected_groups != "<ALL>":
+            self.include_group_names = [x.strip() for x in selected_groups.split(",")]
+        return True
+
+    def _configure_external(self):
+        self.group_match_by_external_id = True
+        return True
+
+    @staticmethod
+    def _is_valid_group_str(group_str: str):
+        return group_str and not re.search(r"[\s#,+ \\<>;]", group_str)
+
+    @staticmethod
+    def _validate_regex(regex_input: str) -> bool:
+        try:
+            re.compile(regex_input)
+            return True
+        except re.error:
+            logger.error(f"{regex_input} is an invalid regular expression")
+            return False

tests/unit/test_install.py

@@ -276,6 +276,7 @@ def not_found(_):
         promtps=MockPrompts(
             {
                 r".*PRO or SERVERLESS SQL warehouse.*": "1",
+                r"Choose how to map the workspace groups.*": "2",
                 r".*": "",
             }
         ),
@@ -411,6 +412,7 @@ def not_found(_):
         promtps=MockPrompts(
             {
                 r".*PRO or SERVERLESS SQL warehouse.*": "1",
+                r"Choose how to map the workspace groups.*": "2",
                 r".*": "",
             }
         ),
@@ -445,6 +447,7 @@ def test_save_config_strip_group_names(ws):
         promtps=MockPrompts(
             {
                 r".*PRO or SERVERLESS SQL warehouse.*": "1",
+                r"Choose how to map the workspace groups.*": "2",  # specify names
                 r".*workspace group names.*": "g1, g2, g99",
                 r".*": "",
             }
@@ -506,6 +509,7 @@ def test_save_config_with_custom_policy(ws):
             {
                 r".*PRO or SERVERLESS SQL warehouse.*": "1",
                 r".*follow a policy.*": "yes",
+                r"Choose how to map the workspace groups.*": "2",
                 r".*Choose a cluster policy.*": "0",
                 r".*": "",
             }
@@ -561,6 +565,7 @@ def test_save_config_with_glue(ws):
         promtps=MockPrompts(
             {
                 r".*PRO or SERVERLESS SQL warehouse.*": "1",
+                r"Choose how to map the workspace groups.*": "2",
                 r".*connect to the external metastore?.*": "yes",
                 r".*Choose a cluster policy.*": "0",
                 r".*": "",

tests/unit/workspace_access/test_groups.py

@@ -7,7 +7,12 @@
 from databricks.sdk.service.iam import ComplexValue, Group, ResourceMeta
 
 from databricks.labs.ucx.framework.parallel import ManyError
-from databricks.labs.ucx.workspace_access.groups import GroupManager, MigratedGroup
+from databricks.labs.ucx.framework.tui import MockPrompts
+from databricks.labs.ucx.workspace_access.groups import (
+    ConfigureGroups,
+    GroupManager,
+    MigratedGroup,
+)
 from tests.unit.framework.mocks import MockBackend
 
 
@@ -679,3 +684,17 @@ def test_snapshot_with_group_matched_by_external_id():
             entitlements='[{"value": "allow-cluster-create"}, {"value": "allow-instance-pool-create"}]',
         )
     ]
+
+
+def test_configure_groups():
+    cg = ConfigureGroups(
+        MockPrompts(
+            {
+                "Backup prefix": "",
+                r"Choose how to map the workspace groups.*": "2",  # specify names
+                r"^Comma-separated list of workspace group names to migrate.*": "foo, bar,  baz",
+            }
+        )
+    )
+    cg.run()
+    assert ["foo", "bar", "baz"] == cg.include_group_names