Add a check for existing inventory database to avoid losing existing, inject installation objects in tests and try fetching existing installation before setting global as default (#1043)

Author: pritishpai

src/databricks/labs/ucx/install.py

@@ -13,7 +13,7 @@
 
 import databricks.sdk.errors
 from databricks.labs.blueprint.entrypoint import get_logger
-from databricks.labs.blueprint.installation import Installation
+from databricks.labs.blueprint.installation import Installation, SerdeError
 from databricks.labs.blueprint.installer import InstallState
 from databricks.labs.blueprint.parallel import ManyError, Threads
 from databricks.labs.blueprint.tui import Prompts
@@ -265,6 +265,8 @@ def _configure_new_installation(self) -> WorkspaceConfig:
         inventory_database = self._prompts.question(
             "Inventory Database stored in hive_metastore", default="ucx", valid_regex=r"^\w+$"
         )
+
+        self._check_inventory_database_exists(inventory_database)
         warehouse_id = self._configure_warehouse()
         configure_groups = ConfigureGroups(self._prompts)
         configure_groups.run()
@@ -331,6 +333,20 @@ def warehouse_type(_):
             warehouse_id = new_warehouse.id
         return warehouse_id
 
+    def _check_inventory_database_exists(self, inventory_database: str):
+        logger.info("Fetching installations...")
+        for installation in Installation.existing(self._ws, self._product_info.product_name()):
+            try:
+                config = installation.load(WorkspaceConfig)
+                if config.inventory_database == inventory_database:
+                    raise AlreadyExists(
+                        f"Inventory database '{inventory_database}' already exists in another installation"
+                    )
+            except NotFound:
+                continue
+            except SerdeError:
+                continue
+
 
 class WorkspaceInstallation:
     def __init__(
@@ -991,6 +1007,9 @@ def validate_and_run(self, step: str):
     logger.setLevel("INFO")
     app = ProductInfo.from_class(WorkspaceConfig)
     workspace_client = WorkspaceClient(product="ucx", product_version=__version__)
-    current = Installation.assume_global(workspace_client, app.product_name())
+    try:
+        current = app.current_installation(workspace_client)
+    except NotFound:
+        current = Installation.assume_global(workspace_client, app.product_name())
     installer = WorkspaceInstaller(Prompts(), current, workspace_client, app)
     installer.run()

tests/integration/test_installation.py

@@ -13,7 +13,7 @@
 from databricks.labs.blueprint.parallel import Threads
 from databricks.labs.blueprint.tui import MockPrompts
 from databricks.labs.blueprint.wheels import ProductInfo
-from databricks.sdk.errors import InvalidParameterValue, NotFound
+from databricks.sdk.errors import AlreadyExists, InvalidParameterValue, NotFound
 from databricks.sdk.retries import retried
 from databricks.sdk.service import compute, sql
 from databricks.sdk.service.iam import PermissionLevel
@@ -38,11 +38,11 @@ def new_installation(ws, sql_backend, env_or_skip, inventory_schema):
 
     def factory(
         config_transform: Callable[[WorkspaceConfig], WorkspaceConfig] | None = None,
+        installation: Installation | None = None,
         product_info: ProductInfo | None = None,
         environ: dict[str, str] | None = None,
         extend_prompts: dict[str, str] | None = None,
-        single_user_install: bool = False,
-        fresh_install: bool = True,
+        inventory_schema_suffix: str = "",
     ):
         if not product_info:
             product_info = ProductInfo.for_testing(WorkspaceConfig)
@@ -58,7 +58,7 @@ def factory(
                 r".*PRO or SERVERLESS SQL warehouse.*": "1",
                 r"Choose how to map the workspace groups.*": "1",
                 r".*connect to the external metastore?.*": "yes",
-                r".*Inventory Database.*": inventory_schema,
+                r".*Inventory Database.*": f"{inventory_schema}{inventory_schema_suffix}",
                 r".*Backup prefix*": renamed_group_prefix,
                 r".*": "",
             }
@@ -75,19 +75,8 @@ def factory(
             ],
         )
 
-        if not fresh_install:
-            installation = product_info.current_installation(ws)
-        elif single_user_install:
-            installation = Installation(
-                ws,
-                product_info.product_name(),
-                install_folder=f"/Users/{ws.current_user.me().user_name}/.{product_info.product_name()}",
-            )
-        else:
-            installation = Installation(
-                ws, product_info.product_name(), install_folder=f"/Applications/{product_info.product_name()}"
-            )
-
+        if not installation:
+            installation = Installation(ws, product_info.product_name())
         installer = WorkspaceInstaller(prompts, installation, ws, product_info, environ)
         workspace_config = installer.configure()
         installation = product_info.current_installation(ws)
@@ -350,31 +339,36 @@ def test_uninstallation(ws, sql_backend, new_installation):
         sql_backend.execute(f"show tables from hive_metastore.{install.config.inventory_database}")
 
 
-def test_fresh_global_installation(new_installation):
+def test_fresh_global_installation(ws, new_installation):
     product_info = ProductInfo.for_testing(WorkspaceConfig)
     global_installation = new_installation(
         product_info=product_info,
+        installation=Installation.assume_global(ws, product_info.product_name()),
     )
     assert global_installation.folder == f"/Applications/{product_info.product_name()}"
     global_installation.uninstall()
 
 
 def test_fresh_user_installation(ws, new_installation):
     product_info = ProductInfo.for_testing(WorkspaceConfig)
-    user_installation = new_installation(product_info=product_info, single_user_install=True)
+    user_installation = new_installation(
+        product_info=product_info,
+        installation=Installation.assume_user_home(ws, product_info.product_name()),
+    )
     assert user_installation.folder == f"/Users/{ws.current_user.me().user_name}/.{product_info.product_name()}"
     user_installation.uninstall()
 
 
-def test_global_installation_on_existing_global_install(new_installation):
+def test_global_installation_on_existing_global_install(ws, new_installation):
     product_info = ProductInfo.for_testing(WorkspaceConfig)
     existing_global_installation = new_installation(
         product_info=product_info,
+        installation=Installation.assume_global(ws, product_info.product_name()),
     )
     assert existing_global_installation.folder == f"/Applications/{product_info.product_name()}"
     reinstall_global = new_installation(
         product_info=product_info,
-        fresh_install=False,
+        installation=Installation.assume_global(ws, product_info.product_name()),
     )
     assert reinstall_global.folder == f"/Applications/{product_info.product_name()}"
     reinstall_global.uninstall()
@@ -385,13 +379,14 @@ def test_user_installation_on_existing_global_install(ws, new_installation):
     product_info = ProductInfo.for_testing(WorkspaceConfig)
     existing_global_installation = new_installation(
         product_info=product_info,
+        installation=Installation.assume_global(ws, product_info.product_name()),
     )
 
     # warning to be thrown by installer if override environment variable present but no confirmation
     with pytest.raises(RuntimeWarning) as err:
         new_installation(
             product_info=product_info,
-            fresh_install=False,
+            installation=Installation.assume_global(ws, product_info.product_name()),
             environ={'UCX_FORCE_INSTALL': 'user'},
             extend_prompts={
                 r".*UCX is already installed on this workspace.*": 'no',
@@ -402,11 +397,12 @@ def test_user_installation_on_existing_global_install(ws, new_installation):
     # successful override with confirmation
     reinstall_user_force = new_installation(
         product_info=product_info,
-        fresh_install=False,
+        installation=Installation.assume_global(ws, product_info.product_name()),
         environ={'UCX_FORCE_INSTALL': 'user'},
         extend_prompts={
             r".*UCX is already installed on this workspace.*": 'yes',
         },
+        inventory_schema_suffix="_reinstall",
     )
     assert reinstall_user_force.folder == f"/Users/{ws.current_user.me().user_name}/.{product_info.product_name()}"
     reinstall_user_force.uninstall()
@@ -416,7 +412,9 @@ def test_user_installation_on_existing_global_install(ws, new_installation):
 def test_global_installation_on_existing_user_install(ws, new_installation):
     # existing installation at user level
     product_info = ProductInfo.for_testing(WorkspaceConfig)
-    existing_user_installation = new_installation(product_info=product_info, single_user_install=True)
+    existing_user_installation = new_installation(
+        product_info=product_info, installation=Installation.assume_user_home(ws, product_info.product_name())
+    )
     assert (
         existing_user_installation.folder == f"/Users/{ws.current_user.me().user_name}/.{product_info.product_name()}"
     )
@@ -425,7 +423,7 @@ def test_global_installation_on_existing_user_install(ws, new_installation):
     with pytest.raises(RuntimeWarning) as err:
         new_installation(
             product_info=product_info,
-            fresh_install=False,
+            installation=Installation.assume_user_home(ws, product_info.product_name()),
             environ={'UCX_FORCE_INSTALL': 'global'},
             extend_prompts={
                 r".*UCX is already installed on this workspace.*": 'no',
@@ -434,15 +432,34 @@ def test_global_installation_on_existing_user_install(ws, new_installation):
     assert err.value.args[0] == "UCX is already installed, but no confirmation"
 
     # not implemented error with confirmation
-
     with pytest.raises(databricks.sdk.errors.NotImplemented) as err:
         new_installation(
             product_info=product_info,
-            fresh_install=False,
+            installation=Installation.assume_user_home(ws, product_info.product_name()),
             environ={'UCX_FORCE_INSTALL': 'global'},
             extend_prompts={
                 r".*UCX is already installed on this workspace.*": 'yes',
             },
         )
     assert err.value.args[0] == "Migration needed. Not implemented yet."
     existing_user_installation.uninstall()
+
+
+def test_check_inventory_database_exists(ws, new_installation):
+    product_info = ProductInfo.for_testing(WorkspaceConfig)
+    install = new_installation(
+        product_info=product_info,
+        installation=Installation.assume_global(ws, product_info.product_name()),
+    )
+    inventory_database = install.config.inventory_database
+
+    with pytest.raises(AlreadyExists) as err:
+        new_installation(
+            product_info=product_info,
+            installation=Installation.assume_global(ws, product_info.product_name()),
+            environ={'UCX_FORCE_INSTALL': 'user'},
+            extend_prompts={
+                r".*UCX is already installed on this workspace.*": 'yes',
+            },
+        )
+    assert err.value.args[0] == f"Inventory database '{inventory_database}' already exists in another installation"

tests/unit/test_install.py

@@ -1,8 +1,10 @@
+import io
 import json
 from datetime import datetime, timedelta
 from unittest.mock import MagicMock, create_autospec, patch
 
 import pytest
+import yaml
 from databricks.labs.blueprint.installation import Installation, MockInstallation
 from databricks.labs.blueprint.installer import InstallState
 from databricks.labs.blueprint.parallel import ManyError
@@ -16,6 +18,7 @@
 from databricks.labs.lsql.backends import MockBackend
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import (  # pylint: disable=redefined-builtin
+    AlreadyExists,
     InvalidParameterValue,
     NotFound,
     NotImplemented,
@@ -84,6 +87,26 @@ def mock_clusters():
 
 @pytest.fixture
 def ws():
+    state = {
+        "/Applications/ucx/config.yml": yaml.dump(
+            {
+                'version': 1,
+                'inventory_database': 'ucx_exists',
+                'connect': {
+                    'host': '...',
+                    'token': '...',
+                },
+            }
+        ),
+    }
+
+    def download(path: str) -> io.StringIO | io.BytesIO:
+        if path not in state:
+            raise NotFound(path)
+        if ".csv" in path:
+            return io.BytesIO(state[path].encode('utf-8'))
+        return io.StringIO(state[path])
+
     workspace_client = create_autospec(WorkspaceClient)
 
     workspace_client.current_user.me = lambda: iam.User(
@@ -107,6 +130,7 @@ def ws():
     workspace_client.cluster_policies.create.return_value = CreatePolicyResponse(policy_id="foo")
     workspace_client.clusters.select_spark_version = lambda latest: "14.2.x-scala2.12"
     workspace_client.clusters.select_node_type = lambda local_disk: "Standard_F4s"
+    workspace_client.workspace.download = download
 
     return workspace_client
 
@@ -503,6 +527,7 @@ def test_create_cluster_policy(ws, mock_installation):
             r".*": "",
         }
     )
+    ws.workspace.get_status = not_found
     install = WorkspaceInstaller(prompts, mock_installation, ws, PRODUCT_INFO)
     install.configure()
     mock_installation.assert_file_written(
@@ -763,7 +788,6 @@ def test_repair_run(ws, mocker, any_prompt, mock_installation_with_jobs):
             state=RunState(result_state=RunResultState.FAILED),
         )
     ]
-    mocker.patch("webbrowser.open")
     ws.jobs.list_runs.return_value = base
     ws.jobs.list_runs.repair_run = None
 
@@ -1087,7 +1111,7 @@ def test_save_config_should_include_databases(ws, mock_installation):
             r".*": "",
         }
     )
-
+    ws.workspace.get_status = not_found
     install = WorkspaceInstaller(prompts, mock_installation, ws, PRODUCT_INFO)
     install.configure()
 
@@ -1190,7 +1214,6 @@ def test_fresh_install(ws, mock_installation):
 
 
 def test_get_existing_installation_global(ws, mock_installation, mocker):
-    mocker.patch("webbrowser.open")
     prompts = MockPrompts(
         {
             r".*PRO or SERVERLESS SQL warehouse.*": "1",
@@ -1221,9 +1244,8 @@ def test_get_existing_installation_global(ws, mock_installation, mocker):
 
     # test for force user install variable without prompts
     second_install = WorkspaceInstaller(prompts, installation, ws, PRODUCT_INFO, force_user_environ)
-    with pytest.raises(RuntimeWarning) as err:
+    with pytest.raises(RuntimeWarning, match='UCX is already installed, but no confirmation'):
         second_install.configure()
-    assert err.value.args[0] == 'UCX is already installed, but no confirmation'
 
     # test for force user install variable with prompts
     prompts = MockPrompts(
@@ -1318,3 +1340,23 @@ def test_databricks_runtime_version_set(ws, mock_installation):
 
     with pytest.raises(SystemExit, match="WorkspaceInstaller is not supposed to be executed in Databricks Runtime"):
         WorkspaceInstaller(prompts, mock_installation, ws, product_info, environ)
+
+
+def test_check_inventory_database_exists(ws, mock_installation):
+    ws.current_user.me().user_name = "foo"
+
+    prompts = MockPrompts(
+        {
+            r".*Inventory Database stored in hive_metastore": "ucx_exists",
+            r".*": "",
+        }
+    )
+
+    installation_type_mock = create_autospec(Installation)
+    installation_type_mock.load.side_effect = NotFound
+
+    installation = Installation(ws, 'ucx')
+    install = WorkspaceInstaller(prompts, installation, ws, PRODUCT_INFO)
+
+    with pytest.raises(AlreadyExists, match="Inventory database 'ucx_exists' already exists in another installation"):
+        install.configure()