Updated Installation to handle install errors

src/databricks/labs/remorph/contexts/application.py

@@ -1,4 +1,3 @@
-# pylint: disable=too-many-public-methods
 import logging
 from functools import cached_property
 
@@ -125,14 +124,10 @@ def workspace_installation(self) -> WorkspaceInstallation:
             self.prompts,
             self.installation,
             self.recon_deployment,
-            self.wheels,
+            self.product_info,
             self.upgrades,
         )
 
     @cached_property
     def upgrades(self):
         return Upgrades(self.product_info, self.installation)
-
-    @cached_property
-    def wheels(self):
-        return self.product_info.wheels(self.workspace_client)

src/databricks/labs/remorph/deployment/installation.py

@@ -1,12 +1,15 @@
 import logging
+from ast import literal_eval
+from pathlib import Path
 
 from databricks.labs.blueprint.installation import Installation
 from databricks.labs.blueprint.tui import Prompts
 from databricks.labs.blueprint.upgrades import Upgrades
-from databricks.labs.blueprint.wheels import WheelsV2
+from databricks.labs.blueprint.wheels import ProductInfo, Version
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import NotFound
-from databricks.sdk.errors.platform import InvalidParameterValue
+from databricks.sdk.mixins.compute import SemVer
+from databricks.sdk.errors.platform import InvalidParameterValue, ResourceDoesNotExist
 
 from databricks.labs.remorph.config import RemorphConfigs
 from databricks.labs.remorph.deployment.recon import ReconDeployment
@@ -21,34 +24,79 @@ def __init__(
         prompts: Prompts,
         installation: Installation,
         recon_deployment: ReconDeployment,
-        wheels: WheelsV2,
+        product_info: ProductInfo,
         upgrades: Upgrades,
     ):
         self._ws = ws
         self._prompts = prompts
         self._installation = installation
         self._recon_deployment = recon_deployment
-        self._wheels = wheels
+        self._product_info = product_info
         self._upgrades = upgrades
 
-    def _apply_upgrades(self):
+    def _get_local_version_file_path(self):
+        user_home = f"{Path(__file__).home()}"
+        return Path(f"{user_home}/.databricks/labs/{self._product_info.product_name()}/state/version.json")
+
+    def _get_local_version_file(self, file_path: Path):
+        data = None
+        with file_path.open("r") as f:
+            data = literal_eval(f.read())
+        assert data, "Unable to read local version file."
+        local_installed_version = data["version"]
         try:
-            self._upgrades.apply(self._ws)
-        except (InvalidParameterValue, NotFound) as err:
-            logger.warning(f"Unable to apply Upgrades due to: {err}")
+            SemVer.parse(local_installed_version)
+        except ValueError:
+            logger.warning(f"{local_installed_version} is not a valid version.")
+            local_installed_version = "v0.3.0"
+        local_installed_date = data["date"]
+        logger.debug(f"Found local installation version: {local_installed_version} {local_installed_date}")
+        return Version(
+            version=local_installed_version,
+            date=local_installed_date,
+            wheel=f"databricks_labs_remorph-{local_installed_version}-py3-none-any.whl",
+        )
+
+    def _get_ws_version(self):
+        try:
+            return self._installation.load(Version)
+        except ResourceDoesNotExist as err:
+            logger.warning(f"Unable to get Workspace Version due to: {err}")
+            return None
+
+    def _apply_upgrades(self):
+        """
+        * If remote version doesn't exist and local version exists:
+           Upload Version file to workspace to handle previous installations.
+        * If remote version or local_version exists, then only apply upgrades.
+        * No need to apply upgrades for fresh installation.
+        """
+        ws_version = self._get_ws_version()
+        local_version_path = self._get_local_version_file_path()
+        local_version = local_version_path.exists()
+        if not ws_version and local_version:
+            self._installation.save(self._get_local_version_file(local_version_path))
+
+        if ws_version or local_version:
+            try:
+                self._upgrades.apply(self._ws)
+                logger.debug("Upgrades applied successfully.")
+            except (InvalidParameterValue, NotFound) as err:
+                logger.warning(f"Unable to apply Upgrades due to: {err}")
 
     def _upload_wheel(self):
-        with self._wheels:
-            wheel_paths = [self._wheels.upload_to_wsfs()]
+        wheels = self._product_info.wheels(self._ws)
+        with wheels:
+            wheel_paths = [wheels.upload_to_wsfs()]
             wheel_paths = [f"/Workspace{wheel}" for wheel in wheel_paths]
             return wheel_paths
 
     def install(self, config: RemorphConfigs):
+        self._apply_upgrades()
         wheel_paths: list[str] = self._upload_wheel()
         if config.reconcile:
             logger.info("Installing Remorph reconcile Metadata components.")
             self._recon_deployment.install(config.reconcile, wheel_paths)
-        self._apply_upgrades()
 
     def uninstall(self, config: RemorphConfigs):
         # This will remove all the Remorph modules

src/databricks/labs/remorph/upgrades/v0.4.0_add_main_table_operation_name_column.py

@@ -157,7 +157,8 @@ def _upgrade_reconcile_metadata_main_table(
 def _upgrade_reconcile_workflow(app_context: ApplicationContext):
     if app_context.recon_config:
         logger.info("Upgrading reconcile workflow")
-        wheel_path = f"/Workspace{app_context.wheels.upload_to_wsfs()}"
+        wheels = app_context.product_info.wheels(app_context.workspace_client)
+        wheel_path = f"/Workspace{wheels.upload_to_wsfs()}"
         app_context.job_deployment.deploy_recon_job(RECON_JOB_NAME, app_context.recon_config, wheel_path)
         logger.debug("Upgraded reconcile workflow")
 

tests/unit/deployment/test_installation.py

@@ -3,7 +3,7 @@
 import pytest
 from databricks.labs.blueprint.installation import MockInstallation, Installation
 from databricks.labs.blueprint.tui import MockPrompts
-from databricks.labs.blueprint.wheels import WheelsV2
+from databricks.labs.blueprint.wheels import WheelsV2, ProductInfo
 from databricks.labs.blueprint.upgrades import Upgrades
 
 from databricks.sdk import WorkspaceClient
@@ -38,7 +38,7 @@ def test_install_all(ws):
     )
     recon_deployment = create_autospec(ReconDeployment)
     installation = create_autospec(Installation)
-    wheels = create_autospec(WheelsV2)
+    product_info = create_autospec(ProductInfo)
     upgrades = create_autospec(Upgrades)
 
     transpile_config = MorphConfig(
@@ -66,15 +66,15 @@ def test_install_all(ws):
         ),
     )
     config = RemorphConfigs(morph=transpile_config, reconcile=reconcile_config)
-    installation = WorkspaceInstallation(ws, prompts, installation, recon_deployment, wheels, upgrades)
+    installation = WorkspaceInstallation(ws, prompts, installation, recon_deployment, product_info, upgrades)
     installation.install(config)
 
 
 def test_no_recon_component_installation(ws):
     prompts = MockPrompts({})
     recon_deployment = create_autospec(ReconDeployment)
     installation = create_autospec(Installation)
-    wheels = create_autospec(WheelsV2)
+    product_info = create_autospec(ProductInfo)
     upgrades = create_autospec(Upgrades)
 
     transpile_config = MorphConfig(
@@ -87,7 +87,7 @@ def test_no_recon_component_installation(ws):
         mode="current",
     )
     config = RemorphConfigs(morph=transpile_config)
-    installation = WorkspaceInstallation(ws, prompts, installation, recon_deployment, wheels, upgrades)
+    installation = WorkspaceInstallation(ws, prompts, installation, recon_deployment, product_info, upgrades)
     installation.install(config)
     recon_deployment.install.assert_not_called()
 
@@ -96,7 +96,7 @@ def test_recon_component_installation(ws):
     recon_deployment = create_autospec(ReconDeployment)
     installation = create_autospec(Installation)
     prompts = MockPrompts({})
-    wheels = create_autospec(WheelsV2)
+    product_info = create_autospec(ProductInfo)
     upgrades = create_autospec(Upgrades)
 
     reconcile_config = ReconcileConfig(
@@ -115,7 +115,7 @@ def test_recon_component_installation(ws):
         ),
     )
     config = RemorphConfigs(reconcile=reconcile_config)
-    installation = WorkspaceInstallation(ws, prompts, installation, recon_deployment, wheels, upgrades)
+    installation = WorkspaceInstallation(ws, prompts, installation, recon_deployment, product_info, upgrades)
     installation.install(config)
     recon_deployment.install.assert_called()