Skip installation if remote and local version is the same, provide prompt to override (#1084)

## Changes
Compares the remote version and local wheel version and prompts for
re-install confirmation if they are the same.

### Linked issues
Closes #1072
Closes #803

### Functionality 

- [ ] added relevant user documentation
- [ ] added new CLI command
- [ ] modified existing command: `databricks labs ucx ...`
- [ ] added a new workflow
- [x] modified existing workflow: `...`
- [ ] added a new table
- [ ] modified existing table: `...`

### Tests
<!-- How is this tested? Please see the checklist below and also
describe any other relevant tests -->

- [ ] manually tested
- [x] added unit tests
- [ ] added integration tests
- [ ] verified on staging environment (screenshot attached)

Author: pritishpai

src/databricks/labs/ucx/install.py

@@ -1,6 +1,7 @@
 import functools
 import logging
 import os
+import re
 import time
 import webbrowser
 from collections.abc import Callable
@@ -13,7 +14,12 @@
 from databricks.labs.blueprint.parallel import ManyError, Threads
 from databricks.labs.blueprint.tui import Prompts
 from databricks.labs.blueprint.upgrades import Upgrades
-from databricks.labs.blueprint.wheels import ProductInfo, WheelsV2, find_project_root
+from databricks.labs.blueprint.wheels import (
+    ProductInfo,
+    Version,
+    WheelsV2,
+    find_project_root,
+)
 from databricks.labs.lsql.backends import SqlBackend, StatementExecutionBackend
 from databricks.labs.lsql.deployment import SchemaDeployer
 from databricks.sdk import WorkspaceClient
@@ -91,6 +97,13 @@ def deploy_schema(sql_backend: SqlBackend, inventory_schema: str):
     deployer.deploy_view("table_estimates", "queries/views/table_estimates.sql")
 
 
+def extract_major_minor(version_string):
+    match = re.search(r'(\d+\.\d+)', version_string)
+    if match:
+        return match.group(1)
+    return None
+
+
 class WorkspaceInstaller:
     def __init__(
         self,
@@ -144,6 +157,20 @@ def run(
                 raise err.errs[0] from None
             raise err
 
+    def _compare_remote_local_versions(self):
+        try:
+            local_version = self._product_info.released_version()
+            remote_version = self._installation.load(Version).version
+            if extract_major_minor(remote_version) == extract_major_minor(local_version):
+                logger.info(f"UCX v{self._product_info.version()} is already installed on this workspace")
+                msg = "Do you want to update the existing installation?"
+                if not self._prompts.confirm(msg):
+                    raise RuntimeWarning(
+                        "UCX workspace remote and local install versions are same and no override is requested. Exiting..."
+                    )
+        except NotFound as err:
+            logger.warning(f"UCX workspace remote version not found: {err}")
+
     def _new_wheel_builder(self):
         return WheelsV2(self._installation, self._product_info)
 
@@ -171,6 +198,7 @@ def _confirm_force_install(self) -> bool:
     def configure(self) -> WorkspaceConfig:
         try:
             config = self._installation.load(WorkspaceConfig)
+            self._compare_remote_local_versions()
             if self._confirm_force_install():
                 return self._configure_new_installation()
             self._apply_upgrades()

tests/integration/test_installation.py

@@ -53,7 +53,6 @@ def factory(
         if not environ:
             environ = {}
         renamed_group_prefix = f"rename-{product_info.product_name()}-"
-
         prompts = MockPrompts(
             {
                 r'Open job overview in your browser.*': 'no',
@@ -377,6 +376,9 @@ def test_global_installation_on_existing_global_install(ws, new_installation):
     reinstall_global, _ = new_installation(
         product_info=product_info,
         installation=Installation.assume_global(ws, product_info.product_name()),
+        extend_prompts={
+            r".*Do you want to update the existing installation?.*": 'yes',
+        },
     )
     assert reinstall_global.folder == f"/Applications/{product_info.product_name()}"
     reinstall_global.uninstall()
@@ -391,16 +393,16 @@ def test_user_installation_on_existing_global_install(ws, new_installation):
     )
 
     # warning to be thrown by installer if override environment variable present but no confirmation
-    with pytest.raises(RuntimeWarning) as err:
+    with pytest.raises(RuntimeWarning, match="UCX is already installed, but no confirmation"):
         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.*": 'no',
+                r".*Do you want to update the existing installation?.*": 'yes',
             },
         )
-    assert err.value.args[0] == "UCX is already installed, but no confirmation"
 
     # successful override with confirmation
     reinstall_user_force, _ = new_installation(
@@ -409,6 +411,7 @@ def test_user_installation_on_existing_global_install(ws, new_installation):
         environ={'UCX_FORCE_INSTALL': 'user'},
         extend_prompts={
             r".*UCX is already installed on this workspace.*": 'yes',
+            r".*Do you want to update the existing installation?.*": 'yes',
         },
         inventory_schema_suffix="_reinstall",
     )
@@ -428,28 +431,28 @@ def test_global_installation_on_existing_user_install(ws, new_installation):
     )
 
     # warning to be thrown by installer if override environment variable present but no confirmation
-    with pytest.raises(RuntimeWarning) as err:
+    with pytest.raises(RuntimeWarning, match="UCX is already installed, but no confirmation"):
         new_installation(
             product_info=product_info,
             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',
+                r".*Do you want to update the existing installation?.*": 'yes',
             },
         )
-    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:
+    with pytest.raises(databricks.sdk.errors.NotImplemented, match="Migration needed. Not implemented yet."):
         new_installation(
             product_info=product_info,
             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',
+                r".*Do you want to update the existing installation?.*": 'yes',
             },
         )
-    assert err.value.args[0] == "Migration needed. Not implemented yet."
     existing_user_installation.uninstall()
 
 
@@ -461,16 +464,18 @@ def test_check_inventory_database_exists(ws, new_installation):
     )
     inventory_database = install.config.inventory_database
 
-    with pytest.raises(AlreadyExists) as err:
+    with pytest.raises(
+        AlreadyExists, match=f"Inventory database '{inventory_database}' already exists in another installation"
+    ):
         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',
+                r".*Do you want to update the existing installation?.*": 'yes',
             },
         )
-    assert err.value.args[0] == f"Inventory database '{inventory_database}' already exists in another installation"
 
 
 @retried(on=[NotFound], timeout=timedelta(minutes=10))
@@ -616,3 +621,20 @@ def test_partitioned_tables(ws, sql_backend, new_installation, inventory_schema,
     assert len(all_tables) >= 2
     assert all_tables[f"{schema.full_name}.partitioned_table"].is_partitioned is True
     assert all_tables[f"{schema.full_name}.non_partitioned_table"].is_partitioned is False
+
+
+def test_compare_remote_local_install_versions(ws, new_installation):
+    product_info = ProductInfo.for_testing(WorkspaceConfig)
+    new_installation(product_info=product_info)
+    with pytest.raises(
+        RuntimeWarning,
+        match="UCX workspace remote and local install versions are same and no override is requested. Exiting...",
+    ):
+        new_installation(product_info=product_info)
+
+    new_installation(
+        product_info=product_info,
+        extend_prompts={
+            r".*Do you want to update the existing installation?.*": 'yes',
+        },
+    )

tests/unit/test_install.py

@@ -55,7 +55,11 @@
 import databricks.labs.ucx.uninstall  # noqa
 from databricks.labs.ucx.config import WorkspaceConfig
 from databricks.labs.ucx.framework.dashboards import DashboardFromFiles
-from databricks.labs.ucx.install import WorkspaceInstallation, WorkspaceInstaller
+from databricks.labs.ucx.install import (
+    WorkspaceInstallation,
+    WorkspaceInstaller,
+    extract_major_minor,
+)
 from databricks.labs.ucx.installer.workflows import WorkflowsInstallation
 
 PRODUCT_INFO = ProductInfo.from_class(WorkspaceConfig)
@@ -1497,3 +1501,69 @@ def test_validate_step(ws, any_prompt, result_state, expected):
     )
 
     assert workflows_installer.validate_step("assessment") == expected
+
+
+def test_are_remote_local_versions_equal(ws, mock_installation, mocker):
+    ws.jobs.run_now = mocker.Mock()
+
+    mocker.patch("webbrowser.open")
+    base_prompts = MockPrompts(
+        {
+            r"Open config file in.*": "yes",
+            r"Open job overview in your browser.*": "yes",
+            r"Do you want to trigger assessment job ?.*": "yes",
+            r"Open assessment Job url that just triggered ?.*": "yes",
+            r".*": "",
+        }
+    )
+
+    product_info = create_autospec(ProductInfo)
+    product_info.released_version.return_value = "0.3.0"
+
+    installation = MockInstallation(
+        {
+            'config.yml': {
+                'inventory_database': 'ucx_user',
+                'connect': {
+                    'host': '...',
+                    'token': '...',
+                },
+            },
+            'version.json': {'version': '0.3.0', 'wheel': '...', 'date': '...'},
+        },
+        is_global=False,
+    )
+
+    install = WorkspaceInstaller(base_prompts, installation, ws, product_info)
+
+    # raises runtime warning when versions match and no override provided
+    with pytest.raises(
+        RuntimeWarning,
+        match="UCX workspace remote and local install versions are same and no override is requested. Exiting...",
+    ):
+        install.configure()
+
+    first_prompts = base_prompts.extend(
+        {
+            r"Do you want to update the existing installation?": "yes",
+        }
+    )
+    install = WorkspaceInstaller(first_prompts, installation, ws, product_info)
+
+    # finishes successfully when versions match and override is provided
+    config = install.configure()
+    assert config.inventory_database == "ucx_user"
+
+    # finishes successfully when versions don't match and no override is provided/needed
+    product_info.released_version.return_value = "0.4.1"
+    install = WorkspaceInstaller(base_prompts, installation, ws, product_info)
+    config = install.configure()
+    assert config.inventory_database == "ucx_user"
+
+
+def test_extract_major_minor_versions():
+    version_string1 = "0.3.123151"
+    version_string2 = "0.17.1232141"
+
+    assert extract_major_minor(version_string1) == "0.3"
+    assert extract_major_minor(version_string2) == "0.17"