diff --git a/src/databricks/labs/ucx/workspace_access/generic.py b/src/databricks/labs/ucx/workspace_access/generic.py index 3eab83acba..d332c9e45c 100644 --- 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 index b370c24023..22ff0f389d 100644 --- 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():