be more defensive

Author: ericvergnaud

src/databricks/labs/ucx/sequencing/sequencing.py

@@ -4,6 +4,7 @@
 from collections.abc import Iterable
 from dataclasses import dataclass, field
 
+from databricks.sdk import WorkspaceClient
 from databricks.sdk.service import jobs
 
 from databricks.labs.ucx.source_code.graph import DependencyGraph
@@ -66,7 +67,8 @@ def find(self, object_type: str, object_id: str) -> MigrationNode | None:
 
 class MigrationSequencer:
 
-    def __init__(self):
+    def __init__(self, ws: WorkspaceClient):
+        self._ws = ws
         self._root = MigrationNode(
             node_id=0, object_type="ROOT", object_id="ROOT", object_name="ROOT", object_owner="NONE"
         )
@@ -83,7 +85,7 @@ def register_workflow_task(self, task: jobs.Task, job: jobs.Job, _graph: Depende
             object_type="TASK",
             object_id=task_id,
             object_name=task.task_key,
-            object_owner=job_node.object_owner, # no task owner so use job one
+            object_owner=job_node.object_owner,  # no task owner so use job one
         )
         job_node.required_steps.append(task_node)
         if task.existing_cluster_id:
@@ -127,14 +129,17 @@ def register_cluster(self, cluster_key: str) -> MigrationNode:
         cluster_node = self._find_node(object_type="CLUSTER", object_id=cluster_key)
         if cluster_node:
             return cluster_node
+        details = self._ws.clusters.get(cluster_key)
+        object_name = details.cluster_name if details and details.cluster_name else cluster_key
+        object_owner = details.creator_user_name if details and details.creator_user_name else "<UNKNOWN>"
         MigrationNode.last_node_id += 1
         cluster_node = MigrationNode(
             node_id=MigrationNode.last_node_id,
             object_type="CLUSTER",
             object_id=cluster_key,
-            object_name=cluster_key,
-            object_owner="NONE",
-        )  # TODO object_owner
+            object_name=object_name,
+            object_owner=object_owner,
+        )
         # TODO register warehouses and policies
         self._root.required_steps.append(cluster_node)
         return cluster_node
@@ -155,6 +160,8 @@ def _deduplicate_steps(steps: Iterable[MigrationStep]) -> Iterable[MigrationStep
         for step in steps:
             existing = best_steps.get(step.step_id, None)
             # keep the step with the highest step number
+            # TODO this possibly affects the step_number of steps that depend on this one
+            # but it's probably OK to not be 100% accurate initially
             if existing and existing.step_number >= step.step_number:
                 continue
             best_steps[step.step_id] = step

tests/unit/sequencing/test_sequencing.py

@@ -1,4 +1,5 @@
 from databricks.sdk.service import jobs
+from databricks.sdk.service.compute import ClusterDetails
 
 from databricks.labs.ucx.sequencing.sequencing import MigrationSequencer
 from databricks.labs.ucx.source_code.base import CurrentSessionState
@@ -7,16 +8,21 @@
 
 
 def test_cluster_from_task_has_children(ws, simple_dependency_resolver, mock_path_lookup):
+    ws.clusters.get.return_value = ClusterDetails(cluster_name="my-cluster", creator_user_name="John Doe")
     task = jobs.Task(task_key="test-task", existing_cluster_id="cluster-123")
     settings = jobs.JobSettings(name="test-job", tasks=[task])
     job = jobs.Job(job_id=1234, settings=settings)
     ws.jobs.get.return_value = job
     dependency = WorkflowTask(ws, task, job)
     graph = DependencyGraph(dependency, None, simple_dependency_resolver, mock_path_lookup, CurrentSessionState())
-    sequencer = MigrationSequencer()
+    sequencer = MigrationSequencer(ws)
     sequencer.register_workflow_task(task, job, graph)
     steps = list(sequencer.generate_steps())
     step = steps[-1]
+    assert step.step_id
     assert step.object_type == "CLUSTER"
     assert step.object_id == "cluster-123"
+    assert step.object_name == "my-cluster"
+    assert step.object_owner == "John Doe"
     assert step.step_number == 3
+    assert len(step.required_step_ids) == 2