Cleanup installer framework and speed up test execution (#711)

Author: nfx

pyproject.toml

@@ -46,6 +46,7 @@ dependencies = [
     "pytest-xdist",
     "pytest-cov>=4.0.0,<5.0.0",
     "pytest-mock>=3.0.0,<4.0.0",
+    "pytest-timeout",
     "black>=23.1.0",
     "ruff>=0.0.243",
     "isort>=2.5.0",
@@ -60,8 +61,8 @@ python="3.10"
 path = ".venv"
 
 [tool.hatch.envs.default.scripts]
-test        = "pytest -n auto --cov src --cov-report=xml tests/unit"
-coverage    = "pytest -n auto --cov src tests/unit --cov-report=html"
+test        = "pytest -n auto --cov src --cov-report=xml --timeout 10 tests/unit"
+coverage    = "pytest -n auto --cov src tests/unit --timeout 10 --cov-report=html"
 integration = "pytest -n 10 --cov src tests/integration"
 fmt         = ["isort .",
                "black .",

src/databricks/labs/ucx/configure.py

@@ -17,7 +17,7 @@ def __init__(self, ws, choice_from_dict):
     def _valid_cluster_id(self, cluster_id: str) -> bool:
         return cluster_id is not None and CLUSTER_ID_LENGTH == len(cluster_id)
 
-    def _configure_override_clusters(self):
+    def configure(self):
         """User may override standard job clusters with interactive clusters"""
         logger.info("Configuring cluster overrides from existing clusters")
 

src/databricks/labs/ucx/framework/tui.py

@@ -1,6 +1,10 @@
+import logging
 import re
+from collections.abc import Callable
 from typing import Any
 
+logger = logging.getLogger(__name__)
+
 
 class Prompts:
     def multi_select_from_dict(self, all_prompt: str, item_prompt: str, choices: dict[str, Any]) -> list[Any]:
@@ -39,6 +43,10 @@ def choice(self, text: str, choices: list[Any], *, max_attempts: int = 10, sort:
         msg = f"cannot get answer within {max_attempts} attempt"
         raise ValueError(msg)
 
+    def confirm(self, text: str, *, max_attempts: int = 10):
+        answer = self.question(text, valid_regex=r"[Yy][Ee][Ss]|[Nn][Oo]", default="no", max_attempts=max_attempts)
+        return answer.lower() == "yes"
+
     def question(
         self,
         text: str,
@@ -47,6 +55,7 @@ def question(
         max_attempts: int = 10,
         valid_number: bool = False,
         valid_regex: str | None = None,
+        validate: Callable[[str], bool] | None = None,
     ) -> str:
         default_help = "" if default is None else f"\033[36m (default: {default})\033[0m"
         prompt = f"\033[1m{text}{default_help}: \033[0m"
@@ -59,6 +68,9 @@ def question(
         while attempt < max_attempts:
             attempt += 1
             res = input(prompt)
+            if res and validate:
+                if not validate(res):
+                    continue
             if res and match_regex:
                 if not match_regex.match(res):
                     print(f"\033[31m[ERROR] Not a '{valid_regex}' match: {res}\033[0m\n")
@@ -74,12 +86,18 @@ def question(
 
 
 class MockPrompts(Prompts):
-    def __init__(self, patterns_to_answers: dict):
-        self._questions_to_answers = {re.compile(k): v for k, v in patterns_to_answers.items()}
+    def __init__(self, patterns_to_answers: dict[str, str]):
+        self._questions_to_answers = sorted(
+            [(re.compile(k), v) for k, v in patterns_to_answers.items()], key=lambda _: len(_[0].pattern), reverse=True
+        )
 
-    def question(self, text: str, **_) -> str:
-        for question, answer in self._questions_to_answers.items():
-            if question.match(text):
-                return answer
+    def question(self, text: str, default: str | None = None, **_) -> str:
+        logger.info(f"Asking prompt: {text}")
+        for question, answer in self._questions_to_answers:
+            if not question.search(text):
+                continue
+            if not answer and default:
+                return default
+            return answer
         mocked = f"not mocked: {text}"
         raise ValueError(mocked)

src/databricks/labs/ucx/framework/wheels.py

@@ -0,0 +1,125 @@
+import datetime
+import logging
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+from contextlib import AbstractContextManager
+from pathlib import Path
+
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.mixins.compute import SemVer
+from databricks.sdk.service.workspace import ImportFormat
+
+logger = logging.getLogger(__name__)
+
+
+class Wheels(AbstractContextManager):
+    def __init__(self, ws: WorkspaceClient, install_folder: str, released_version: str):
+        self._ws = ws
+        self._this_file = Path(__file__)
+        self._install_folder = install_folder
+        self._released_version = released_version
+
+    def version(self):
+        if hasattr(self, "__version"):
+            return self.__version
+        project_root = find_project_root()
+        if not (project_root / ".git/config").exists():
+            # normal install, downloaded releases won't have the .git folder
+            return self._released_version
+        try:
+            out = subprocess.run(["git", "describe", "--tags"], stdout=subprocess.PIPE, check=True)  # noqa S607
+            git_detached_version = out.stdout.decode("utf8")
+            dv = SemVer.parse(git_detached_version)
+            datestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+            # new commits on main branch since the last tag
+            new_commits = dv.pre_release.split("-")[0] if dv.pre_release else None
+            # show that it's a version different from the released one in stats
+            bump_patch = dv.patch + 1
+            # create something that is both https://semver.org and https://peps.python.org/pep-0440/
+            semver_and_pep0440 = f"{dv.major}.{dv.minor}.{bump_patch}+{new_commits}{datestamp}"
+            # validate the semver
+            SemVer.parse(semver_and_pep0440)
+            self.__version = semver_and_pep0440
+            return semver_and_pep0440
+        except Exception as err:
+            msg = (
+                f"Cannot determine unreleased version. Please report this error "
+                f"message that you see on https://github.com/databrickslabs/ucx/issues/new. "
+                f"Meanwhile, download, unpack, and install the latest released version from "
+                f"https://github.com/databrickslabs/ucx/releases. Original error is: {err!s}"
+            )
+            raise OSError(msg) from None
+
+    def __enter__(self) -> "Wheels":
+        self._tmp_dir = tempfile.TemporaryDirectory()
+        self._local_wheel = self._build_wheel(self._tmp_dir.name)
+        self._remote_wheel = f"{self._install_folder}/wheels/{self._local_wheel.name}"
+        self._remote_dirname = os.path.dirname(self._remote_wheel)
+        return self
+
+    def __exit__(self, __exc_type, __exc_value, __traceback):
+        self._tmp_dir.cleanup()
+
+    def upload_to_dbfs(self) -> str:
+        with self._local_wheel.open("rb") as f:
+            self._ws.dbfs.mkdirs(self._remote_dirname)
+            logger.info(f"Uploading wheel to dbfs:{self._remote_wheel}")
+            self._ws.dbfs.upload(self._remote_wheel, f, overwrite=True)
+        return self._remote_wheel
+
+    def upload_to_wsfs(self) -> str:
+        with self._local_wheel.open("rb") as f:
+            self._ws.workspace.mkdirs(self._remote_dirname)
+            logger.info(f"Uploading wheel to /Workspace{self._remote_wheel}")
+            self._ws.workspace.upload(self._remote_wheel, f, overwrite=True, format=ImportFormat.AUTO)
+        return self._remote_wheel
+
+    def _build_wheel(self, tmp_dir: str, *, verbose: bool = False):
+        """Helper to build the wheel package"""
+        stdout = subprocess.STDOUT
+        stderr = subprocess.STDOUT
+        if not verbose:
+            stdout = subprocess.DEVNULL
+            stderr = subprocess.DEVNULL
+        project_root = find_project_root()
+        is_non_released_version = "+" in self.version()
+        if (project_root / ".git" / "config").exists() and is_non_released_version:
+            tmp_dir_path = Path(tmp_dir) / "working-copy"
+            # copy everything to a temporary directory
+            shutil.copytree(project_root, tmp_dir_path)
+            # and override the version file
+            # TODO: make it configurable
+            version_file = tmp_dir_path / "src/databricks/labs/ucx/__about__.py"
+            with version_file.open("w") as f:
+                f.write(f'__version__ = "{self.version()}"')
+            # working copy becomes project root for building a wheel
+            project_root = tmp_dir_path
+        logger.debug(f"Building wheel for {project_root} in {tmp_dir}")
+        subprocess.run(
+            [sys.executable, "-m", "pip", "wheel", "--no-deps", "--wheel-dir", tmp_dir, project_root.as_posix()],
+            check=True,
+            stdout=stdout,
+            stderr=stderr,
+        )
+        # get wheel name as first file in the temp directory
+        return next(Path(tmp_dir).glob("*.whl"))
+
+
+def find_project_root() -> Path:
+    def _find_dir_with_leaf(folder: Path, leaf: str) -> Path | None:
+        root = folder.root
+        while str(folder.absolute()) != root:
+            if (folder / leaf).exists():
+                return folder
+            folder = folder.parent
+        return None
+
+    for leaf in ["pyproject.toml", "setup.py"]:
+        root = _find_dir_with_leaf(Path(__file__), leaf)
+        if root is not None:
+            return root
+    msg = "Cannot find project root"
+    raise NotADirectoryError(msg)