From 124656aab2323b5722ae6bf270a2748c7c54de83 Mon Sep 17 00:00:00 2001 From: Lars George Date: Tue, 19 Sep 2023 10:50:53 +0200 Subject: [PATCH 1/2] Allow empty group filter in install.py (#216) Also fixes stripping whitespaces around groups when enters like `g1, g2, g3` etc. Previously the code expected to get no whitespaces. --- src/databricks/labs/ucx/install.py | 29 ++++++++++++++++++++--------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/src/databricks/labs/ucx/install.py b/src/databricks/labs/ucx/install.py index 149dc4d176..f04afeb1d8 100644 --- a/src/databricks/labs/ucx/install.py +++ b/src/databricks/labs/ucx/install.py @@ -110,15 +110,26 @@ def _configure(self): raise err logger.info("Please answer a couple of questions to configure Unity Catalog migration") + inventory_database = self._question("Inventory Database", default="ucx") + selected_groups = self._question( + "Comma-separated list of workspace group names to migrate (empty means all)", default="" + ) + backup_group_prefix = self._question("Backup prefix", default="db-temp-") + log_level = self._question("Log level", default="INFO") + num_threads = int(self._question("Number of threads", default="8")) + groups_config_args = { + "backup_group_prefix": backup_group_prefix, + } + if selected_groups != "": + groups_config_args["selected"] = [x.strip() for x in selected_groups.split(",")] + else: + groups_config_args["auto"] = True self._config = MigrationConfig( - inventory_database=self._question("Inventory Database", default="ucx"), - groups=GroupsConfig( - selected=self._question("Comma-separated list of workspace group names to migrate").split(","), - backup_group_prefix=self._question("Backup prefix", default="db-temp-"), - ), + inventory_database=inventory_database, + groups=GroupsConfig(**groups_config_args), tacl=TaclConfig(auto=True), - log_level=self._question("Log level", default="INFO"), - num_threads=int(self._question("Number of threads", default="8")), + log_level=log_level, + num_threads=num_threads, ) self._write_config() @@ -182,8 +193,8 @@ def _create_readme(self): self._ws.workspace.upload(path, intro.encode("utf8"), overwrite=True) url = self._notebook_link(path) logger.info(f"Created README notebook with job overview: {url}") - msg = "Type 'yes' to open job overview in README notebook in your home directory" - if self._prompts and self._question(msg) == "yes": + msg = "Open job overview in README notebook in your home directory" + if self._prompts and self._question(msg, default="yes") == "yes": webbrowser.open(url) def _create_debug(self, remote_wheel: str): From 4d662663a2de44a6e80e9763b34196acbddfff9c Mon Sep 17 00:00:00 2001 From: Lars George Date: Tue, 19 Sep 2023 11:33:58 +0200 Subject: [PATCH 2/2] Improve unit test --- tests/unit/test_install.py | 94 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 94 insertions(+) diff --git a/tests/unit/test_install.py b/tests/unit/test_install.py index d05efaefdb..94c1283be6 100644 --- a/tests/unit/test_install.py +++ b/tests/unit/test_install.py @@ -2,6 +2,7 @@ import os.path from pathlib import Path +import pytest import yaml from databricks.sdk.core import DatabricksError from databricks.sdk.service import iam @@ -49,6 +50,99 @@ def not_found(_): ) +def test_save_config_with_error(mocker): + def not_found(_): + raise DatabricksError(error_code="RAISED_FOR_TESTING") + + mocker.patch("builtins.input", return_value="42") + ws = mocker.Mock() + ws.current_user.me = lambda: iam.User(user_name="me@example.com", groups=[iam.ComplexValue(display="admins")]) + ws.config.host = "https://foo" + ws.workspace.get_status = not_found + + install = Installer(ws) + with pytest.raises(DatabricksError) as e_info: + install._configure() + assert str(e_info.value.error_code) == "RAISED_FOR_TESTING" + + +def test_save_config_auto_groups(mocker): + def not_found(_): + raise DatabricksError(error_code="RESOURCE_DOES_NOT_EXIST") + + def mock_question(text: str, *, default: str | None = None) -> str: + if "workspace group names" in text: + return "" + else: + return "42" + + mocker.patch("builtins.input", return_value="42") + ws = mocker.Mock() + ws.current_user.me = lambda: iam.User(user_name="me@example.com", groups=[iam.ComplexValue(display="admins")]) + ws.config.host = "https://foo" + ws.workspace.get_status = not_found + + install = Installer(ws) + install._question = mock_question + install._configure() + + ws.workspace.upload.assert_called_with( + "/Users/me@example.com/.ucx/config.yml", + b"""groups: + auto: true + backup_group_prefix: '42' +inventory_database: '42' +log_level: '42' +num_threads: 42 +tacl: + auto: true +version: 1 +workspace_start_path: / +""", + format=ImportFormat.AUTO, + ) + + +def test_save_config_strip_group_names(mocker): + def not_found(_): + raise DatabricksError(error_code="RESOURCE_DOES_NOT_EXIST") + + def mock_question(text: str, *, default: str | None = None) -> str: + if "workspace group names" in text: + return "g1, g2, g99" + else: + return "42" + + mocker.patch("builtins.input", return_value="42") + ws = mocker.Mock() + ws.current_user.me = lambda: iam.User(user_name="me@example.com", groups=[iam.ComplexValue(display="admins")]) + ws.config.host = "https://foo" + ws.workspace.get_status = not_found + + install = Installer(ws) + install._question = mock_question + install._configure() + + ws.workspace.upload.assert_called_with( + "/Users/me@example.com/.ucx/config.yml", + b"""groups: + backup_group_prefix: '42' + selected: + - g1 + - g2 + - g99 +inventory_database: '42' +log_level: '42' +num_threads: 42 +tacl: + auto: true +version: 1 +workspace_start_path: / +""", + format=ImportFormat.AUTO, + ) + + def test_main_with_existing_conf_does_not_recreate_config(mocker): mocker.patch("builtins.input", return_value="yes") webbrowser_open = mocker.patch("webbrowser.open")