Introduced Safer access to workspace objects' properties. (#530)

Closes #521
databrickslabs · Nov 9, 2023 · fb769c4 · fb769c4
1 parent 3b729d1
commit fb769c4
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 10 deletions.
diff --git a/src/databricks/labs/ucx/workspace_access/generic.py b/src/databricks/labs/ucx/workspace_access/generic.py
@@ -33,9 +33,9 @@ class GenericPermissionsInfo:
 
 @dataclass
 class WorkspaceObjectInfo:
-    object_type: str
-    object_id: str
     path: str
+    object_type: str = None
+    object_id: str = None
     language: str = None
 
 
@@ -299,9 +299,9 @@ def _crawl(self) -> list[WorkspaceObjectInfo]:
                 continue
             raw = obj.as_dict()
             yield WorkspaceObjectInfo(
-                object_type=raw["object_type"],
-                object_id=str(raw["object_id"]),
-                path=raw["path"],
+                object_type=raw.get("object_type", None),
+                object_id=str(raw.get("object_id", None)),
+                path=raw.get("path", None),
                 language=raw.get("language", None),
             )
 
@@ -310,7 +310,9 @@ def snapshot(self) -> list[WorkspaceObjectInfo]:
 
     def _try_fetch(self) -> list[WorkspaceObjectInfo]:
         for row in self._fetch(f"SELECT * FROM {self._schema}.{self._table}"):
-            yield WorkspaceObjectInfo(*row)
+            yield WorkspaceObjectInfo(
+                path=row["path"], object_type=row["object_type"], object_id=row["object_id"], language=row["language"]
+            )
 
     def object_types(self) -> set[str]:
         return {"notebooks", "directories", "repos", "files"}

diff --git a/tests/unit/workspace_access/test_generic.py b/tests/unit/workspace_access/test_generic.py
@@ -531,7 +531,7 @@ def test_workspaceobject_try_fetch():
     result_set = list(crawler._try_fetch())
 
     assert len(result_set) == 2
-    assert result_set[0] == WorkspaceObjectInfo("NOTEBOOK", 123, "/rootobj/notebook1", "PYTHON")
+    assert result_set[0] == WorkspaceObjectInfo("/rootobj/notebook1", "NOTEBOOK", 123, "PYTHON")
 
 
 def test_workspaceobject_crawl():
@@ -560,7 +560,7 @@ def test_workspaceobject_crawl():
         result_set = list(crawler)
 
     assert len(result_set) == 2
-    assert result_set[0] == WorkspaceObjectInfo("NOTEBOOK", "123", "/rootobj/notebook1", "PYTHON")
+    assert result_set[0] == WorkspaceObjectInfo("/rootobj/notebook1", "NOTEBOOK", "123", "PYTHON")
 
 
 def test_workspaceobject_withexperiment_crawl():
@@ -587,7 +587,7 @@ def test_workspaceobject_withexperiment_crawl():
         result_set = list(crawler)
 
     assert len(result_set) == 1
-    assert result_set[0] == WorkspaceObjectInfo("NOTEBOOK", "123", "/rootobj/notebook1", "PYTHON")
+    assert result_set[0] == WorkspaceObjectInfo("/rootobj/notebook1", "NOTEBOOK", "123", "PYTHON")
 
 
 def test_workspace_snapshot():
@@ -613,7 +613,7 @@ def test_workspace_snapshot():
     result_set = crawler.snapshot()
 
     assert len(result_set) == 2
-    assert result_set[0] == WorkspaceObjectInfo("NOTEBOOK", "123", "/rootobj/notebook1", "PYTHON")
+    assert result_set[0] == WorkspaceObjectInfo("/rootobj/notebook1", "NOTEBOOK", "123", "PYTHON")
 
 
 def test_eligibles_assets_with_owner_should_be_accepted():