Fixing installer not being able to migrate from V1 to V2 (#596)

- [x] Needs to verify if "backup_group_prefix" is properly migrated
- [x] Needs to verify if "selected_groups" is properly migrated
- [x] Needs to verify if "auto:True" is properly migrated
- [ ] Needs to verify if previous config and jobs are overriden by V2
config
- [x] Needs to work on a fresh installation
- [x] Needs to work when being re-run

Author: william-conti

src/databricks/labs/ucx/config.py

@@ -206,8 +206,11 @@ def _migrate_from_v1(cls, raw: dict):
         if stored_version == _CONFIG_VERSION:
             return raw
         if stored_version == 1:
-            raw["include_group_names"] = raw.get("groups", {"selected": []})["selected"]
+            raw["include_group_names"] = (
+                raw.get("groups", {"selected": []})["selected"] if "selected" in raw["groups"] else None
+            )
             raw["renamed_group_prefix"] = raw.get("groups", {"backup_group_prefix": "db-temp-"})["backup_group_prefix"]
+            raw.pop("groups", None)
             return raw
         msg = f"Unknown config version: {stored_version}"
         raise ValueError(msg)

src/databricks/labs/ucx/install.py

@@ -16,7 +16,7 @@
 import yaml
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.core import DatabricksError
-from databricks.sdk.errors import OperationFailed
+from databricks.sdk.errors import NotFound, OperationFailed
 from databricks.sdk.mixins.compute import SemVer
 from databricks.sdk.service import compute, jobs
 from databricks.sdk.service.sql import EndpointInfoWarehouseType, SpotInstancePolicy
@@ -146,6 +146,7 @@ def __init__(
 
     def run(self):
         logger.info(f"Installing UCX v{self._version}")
+
         self._configure()
         self._run_configured()
 
@@ -215,7 +216,7 @@ def run_for_config(
         workspace_installer = WorkspaceInstaller(ws, prefix=prefix, promtps=False, sql_backend=sql_backend)
         logger.info(f"Installing UCX v{workspace_installer._version} on {ws.config.host}")
         workspace_installer._config = config
-        workspace_installer._write_config()
+        workspace_installer._write_config(overwrite=False)
         workspace_installer._override_clusters = override_clusters
         # TODO: rather introduce a method `is_configured`, as we may want to reconfigure workspaces for some reason
         workspace_installer._run_configured()
@@ -316,6 +317,10 @@ def _current_config(self):
             self._config = WorkspaceConfig.from_bytes(f.read())
         return self._config
 
+    def _raw_previous_config(self):
+        with self._ws.workspace.download(self.config_file) as f:
+            return str(f.read())
+
     def _name(self, name: str) -> str:
         return f"[{self._prefix.upper()}] {name}"
 
@@ -337,13 +342,17 @@ def _configure_inventory_database(self):
     def _configure(self):
         ws_file_url = self.notebook_link(self.config_file)
         try:
-            self._ws.workspace.get_status(self.config_file)
-            logger.info(f"UCX is already configured. See {ws_file_url}")
-            return
-        except DatabricksError as err:
+            if "version: 1" in self._raw_previous_config():
+                logger.info("old version detected, attempting to migrate to new config")
+                self._config = self._current_config
+                self._write_config(overwrite=True)
+                return
+            elif "version: 2" in self._raw_previous_config():
+                logger.info(f"UCX is already configured. See {ws_file_url}")
+                return
+        except NotFound as err:
             if err.error_code != "RESOURCE_DOES_NOT_EXIST":
                 raise err
-
         logger.info("Please answer a couple of questions to configure Unity Catalog migration")
         inventory_database = self._configure_inventory_database()
 
@@ -418,12 +427,12 @@ def warehouse_type(_):
             spark_conf=spark_conf_dict,
         )
 
-        self._write_config()
+        self._write_config(overwrite=False)
         msg = "Open config file in the browser and continue installing?"
         if self._prompts and self._question(msg, default="yes") == "yes":
             webbrowser.open(ws_file_url)
 
-    def _write_config(self):
+    def _write_config(self, overwrite):
         try:
             self._ws.workspace.get_status(self._install_folder)
         except DatabricksError as err:
@@ -434,7 +443,7 @@ def _write_config(self):
 
         config_bytes = yaml.dump(self._config.as_dict()).encode("utf8")
         logger.info(f"Creating configuration file: {self.config_file}")
-        self._ws.workspace.upload(self.config_file, config_bytes, format=ImportFormat.AUTO)
+        self._ws.workspace.upload(self.config_file, config_bytes, format=ImportFormat.AUTO, overwrite=overwrite)
 
     def _create_jobs(self):
         if not self._state.jobs:

tests/unit/test_install.py

@@ -7,7 +7,7 @@
 import pytest
 import yaml
 from databricks.sdk.core import DatabricksError
-from databricks.sdk.errors import OperationFailed
+from databricks.sdk.errors import NotFound, OperationFailed
 from databricks.sdk.service import iam, jobs
 from databricks.sdk.service.compute import (
     GlobalInitScriptDetails,
@@ -41,8 +41,6 @@ def ws(mocker):
     ws.current_user.me = lambda: iam.User(user_name="[email protected]", groups=[iam.ComplexValue(display="admins")])
     ws.config.host = "https://foo"
     ws.config.is_aws = True
-    config_bytes = yaml.dump(WorkspaceConfig(inventory_database="a").as_dict()).encode("utf8")
-    ws.workspace.download = lambda _: io.BytesIO(config_bytes)
     ws.workspace.get_status = lambda _: ObjectInfo(object_id=123)
     ws.data_sources.list = lambda: [DataSource(id="bcd", warehouse_id="abc")]
     ws.warehouses.list = lambda **_: [EndpointInfo(id="abc", warehouse_type=EndpointInfoWarehouseType.PRO)]
@@ -55,6 +53,8 @@ def ws(mocker):
 
 
 def test_replace_clusters_for_integration_tests(ws):
+    config_bytes = yaml.dump(WorkspaceConfig(inventory_database="a").as_dict()).encode("utf8")
+    ws.workspace.download = lambda _: io.BytesIO(config_bytes)
     return_value = WorkspaceInstaller.run_for_config(
         ws,
         WorkspaceConfig(inventory_database="a"),
@@ -148,22 +148,127 @@ def not_found(_):
 workspace_start_path: /
 """,
         format=ImportFormat.AUTO,
+        overwrite=False,
+    )
+
+
+def test_migrate_from_v1(ws, mocker):
+    mocker.patch("builtins.input", return_value="yes")
+    mock_file = MagicMock()
+    mocker.patch("builtins.open", return_value=mock_file)
+    mocker.patch("base64.b64encode")
+
+    ws.current_user.me = lambda: iam.User(user_name="[email protected]", groups=[iam.ComplexValue(display="admins")])
+    ws.config.host = "https://foo"
+    ws.config.is_aws = True
+    config_bytes = b"""default_catalog: ucx_default
+groups:
+  auto: true
+  backup_group_prefix: db-temp-
+inventory_database: ucx
+log_level: INFO
+num_threads: 8
+version: 1
+warehouse_id: abc
+workspace_start_path: /
+    """
+    ws.workspace.download = lambda _: io.BytesIO(config_bytes)
+    ws.workspace.get_status = lambda _: ObjectInfo(object_id=123)
+    mocker.patch("builtins.input", return_value="42")
+
+    ws.warehouses.list = lambda **_: [
+        EndpointInfo(id="abc", warehouse_type=EndpointInfoWarehouseType.PRO, state=State.RUNNING)
+    ]
+    ws.cluster_policies.list = lambda: []
+
+    install = WorkspaceInstaller(ws)
+    install._choice = lambda _1, _2: "None (abc, PRO, RUNNING)"
+    install._configure()
+
+    ws.workspace.upload.assert_called_with(
+        "/Users/[email protected]/.ucx/config.yml",
+        b"""default_catalog: ucx_default
+inventory_database: ucx
+log_level: INFO
+num_threads: 8
+renamed_group_prefix: db-temp-
+version: 2
+warehouse_id: abc
+workspace_start_path: /
+""",
+        format=ImportFormat.AUTO,
+        overwrite=True,
+    )
+
+
+def test_migrate_from_v1_selected_groups(ws, mocker):
+    mocker.patch("builtins.input", return_value="yes")
+    mock_file = MagicMock()
+    mocker.patch("builtins.open", return_value=mock_file)
+    mocker.patch("base64.b64encode")
+
+    ws.current_user.me = lambda: iam.User(user_name="[email protected]", groups=[iam.ComplexValue(display="admins")])
+    ws.config.host = "https://foo"
+    ws.config.is_aws = True
+    config_bytes = b"""default_catalog: ucx_default
+groups:
+  backup_group_prefix: 'backup_baguette_prefix'
+  selected:
+  - '42'
+  - '100'
+inventory_database: '42'
+log_level: '42'
+num_threads: 42
+version: 1
+warehouse_id: abc
+workspace_start_path: /
+    """
+    ws.workspace.download = lambda _: io.BytesIO(config_bytes)
+    ws.workspace.get_status = lambda _: ObjectInfo(object_id=123)
+    mocker.patch("builtins.input", return_value="42")
+
+    ws.warehouses.list = lambda **_: [
+        EndpointInfo(id="abc", warehouse_type=EndpointInfoWarehouseType.PRO, state=State.RUNNING)
+    ]
+    ws.cluster_policies.list = lambda: []
+
+    install = WorkspaceInstaller(ws)
+    install._choice = lambda _1, _2: "None (abc, PRO, RUNNING)"
+    install._configure()
+
+    ws.workspace.upload.assert_called_with(
+        "/Users/[email protected]/.ucx/config.yml",
+        b"""default_catalog: ucx_default
+include_group_names:
+- '42'
+- '100'
+inventory_database: '42'
+log_level: '42'
+num_threads: 42
+renamed_group_prefix: backup_baguette_prefix
+version: 2
+warehouse_id: abc
+workspace_start_path: /
+""",
+        format=ImportFormat.AUTO,
+        overwrite=True,
     )
 
 
 def test_save_config_with_error(ws, mocker):
     def not_found(_):
-        raise DatabricksError(error_code="RAISED_FOR_TESTING")
+        raise NotFound(message="File not found")
 
     mocker.patch("builtins.input", return_value="42")
 
-    ws.workspace.get_status = not_found
+    ws.workspace.download = not_found
     ws.cluster_policies.list = lambda: []
 
     install = WorkspaceInstaller(ws)
-    with pytest.raises(DatabricksError) as e_info:
+    with pytest.raises(NotFound) as e_info:
         install._configure()
-    assert str(e_info.value.error_code) == "RAISED_FOR_TESTING"
+
+    assert str(e_info.value.args[0]) == "File not found"
 
 
 def test_save_config_auto_groups(ws, mocker):
@@ -201,6 +306,7 @@ def mock_question(text: str, *, default: str | None = None) -> str:
 workspace_start_path: /
 """,
         format=ImportFormat.AUTO,
+        overwrite=False,
     )
 
 
@@ -243,6 +349,7 @@ def mock_question(text: str, *, default: str | None = None) -> str:
 workspace_start_path: /
 """,
         format=ImportFormat.AUTO,
+        overwrite=False,
     )
 
 
@@ -307,6 +414,7 @@ def mock_choice_from_dict(text: str, choices: dict[str, Any]) -> Any:
 workspace_start_path: /
 """,
         format=ImportFormat.AUTO,
+        overwrite=False,
     )
 
 
@@ -330,7 +438,7 @@ def test_main_with_existing_conf_does_not_recreate_config(ws, mocker):
 renamed_group_prefix: '42'
 spark_conf:
   spark.databricks.hive.metastore.glueCatalog.enabled: 'true'
-version: 1
+version: 2
 warehouse_id: abc
 workspace_start_path: /
 """