diff --git a/src/databricks/labs/ucx/framework/tasks.py b/src/databricks/labs/ucx/framework/tasks.py index ff5d6270f6..3dc4c758fb 100644 --- a/src/databricks/labs/ucx/framework/tasks.py +++ b/src/databricks/labs/ucx/framework/tasks.py @@ -12,6 +12,7 @@ @dataclass class Task: + task_id: int workflow: str name: str doc: str @@ -54,6 +55,7 @@ def wrapper(*args, **kwargs): raise SyntaxError(msg) _TASKS[func.__name__] = Task( + task_id=len(_TASKS), workflow=workflow, name=func.__name__, doc=func.__doc__, diff --git a/src/databricks/labs/ucx/install.py b/src/databricks/labs/ucx/install.py index 0eccf29c0e..a984bc71c8 100644 --- a/src/databricks/labs/ucx/install.py +++ b/src/databricks/labs/ucx/install.py @@ -263,6 +263,18 @@ def _create_jobs(self): self._create_readme() self._create_debug(remote_wheel) + @staticmethod + def _sorted_tasks() -> list[Task]: + return sorted(_TASKS.values(), key=lambda x: x.task_id) + + @classmethod + def _step_list(cls) -> list[str]: + step_list = [] + for task in cls._sorted_tasks(): + if task.workflow not in step_list: + step_list.append(task.workflow) + return step_list + def _create_readme(self): md = [ "# UCX - The Unity Catalog Migration Assistant", @@ -270,14 +282,18 @@ def _create_readme(self): "Here are the URL and descriptions of jobs that trigger's various stages of migration.", "All jobs are defined with necessary cluster configurations and DBR versions.", ] - for step_name, job_id in self._deployed_steps.items(): + for step_name in self._step_list(): + if step_name not in self._deployed_steps: + logger.warning(f"Skipping step '{step_name}' since it was not deployed.") + continue + job_id = self._deployed_steps[step_name] dashboard_link = "" if step_name in self._dashboards: dashboard_link = f"{self._ws.config.host}/sql/dashboards/{self._dashboards[step_name]}" dashboard_link = f" (see [{step_name} dashboard]({dashboard_link}) after finish)" job_link = f"[{self._name(step_name)}]({self._ws.config.host}#job/{job_id})" md.append(f"## {job_link}{dashboard_link}\n") - for t in _TASKS.values(): + for t in self._sorted_tasks(): if t.workflow != step_name: continue doc = re.sub(r"\s+", " ", t.doc) diff --git a/tests/unit/test_install.py b/tests/unit/test_install.py index f5960280e7..5010458234 100644 --- a/tests/unit/test_install.py +++ b/tests/unit/test_install.py @@ -220,3 +220,56 @@ def test_choices_happy(mocker): mocker.patch("builtins.input", return_value="1") res = install._choice("foo", ["a", "b"]) assert "b" == res + + +def test_step_list(mocker): + ws = mocker.Mock() + from databricks.labs.ucx.framework.tasks import Task + + tasks = [ + Task(task_id=0, workflow="wl_1", name="n3", doc="d3", fn=lambda: None), + Task(task_id=1, workflow="wl_2", name="n2", doc="d2", fn=lambda: None), + Task(task_id=2, workflow="wl_1", name="n1", doc="d1", fn=lambda: None), + ] + + with mocker.patch.object(WorkspaceInstaller, attribute="_sorted_tasks", return_value=tasks): + install = WorkspaceInstaller(ws) + steps = install._step_list() + assert len(steps) == 2 + assert steps[0] == "wl_1" and steps[1] == "wl_2" + + +def test_create_readme(mocker): + mocker.patch("builtins.input", return_value="yes") + webbrowser_open = mocker.patch("webbrowser.open") + ws = mocker.Mock() + + ws.current_user.me = lambda: iam.User(user_name="me@example.com", groups=[iam.ComplexValue(display="admins")]) + ws.config.host = "https://foo" + config_bytes = yaml.dump(WorkspaceConfig(inventory_database="a", groups=GroupsConfig(auto=True)).as_dict()).encode( + "utf8" + ) + ws.workspace.download = lambda _: io.BytesIO(config_bytes) + + from databricks.labs.ucx.framework.tasks import Task + + tasks = [ + Task(task_id=0, workflow="wl_1", name="n3", doc="d3", fn=lambda: None), + Task(task_id=1, workflow="wl_2", name="n2", doc="d2", fn=lambda: None), + Task(task_id=2, workflow="wl_1", name="n1", doc="d1", fn=lambda: None), + ] + + with mocker.patch.object(WorkspaceInstaller, attribute="_sorted_tasks", return_value=tasks): + install = WorkspaceInstaller(ws) + install._deployed_steps = {"wl_1": 1, "wl_2": 2} + install._create_readme() + + webbrowser_open.assert_called_with("https://foo/#workspace/Users/me@example.com/.ucx/README.py") + + _, args, kwargs = ws.mock_calls[0] + assert args[0] == "/Users/me@example.com/.ucx/README.py" + + import re + + p = re.compile(".*wl_1.*n3.*n1.*wl_2.*n2.*") + assert p.match(str(args[1]))