Skip to content

Commit

Permalink
Keep order of tasks in the README notebook (#286)
Browse files Browse the repository at this point in the history
Resolves #285
  • Loading branch information
larsgeorge-db committed Sep 26, 2023
1 parent cd033d5 commit b1cd014
Show file tree
Hide file tree
Showing 3 changed files with 73 additions and 2 deletions.
2 changes: 2 additions & 0 deletions src/databricks/labs/ucx/framework/tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@

@dataclass
class Task:
task_id: int
workflow: str
name: str
doc: str
Expand Down Expand Up @@ -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__,
Expand Down
20 changes: 18 additions & 2 deletions src/databricks/labs/ucx/install.py
Original file line number Diff line number Diff line change
Expand Up @@ -263,21 +263,37 @@ 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",
f'To troubleshoot, see [debug notebook]({self._notebook_link(f"{self._install_folder}/DEBUG.py")}).\n',
"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)
Expand Down
53 changes: 53 additions & 0 deletions tests/unit/test_install.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]))

0 comments on commit b1cd014

Please sign in to comment.