From 9ca2d82608cb5c5b23b8c2ae7efb369d9f761054 Mon Sep 17 00:00:00 2001 From: vuong-nguyen <44292934+nkvuong@users.noreply.github.com> Date: Fri, 3 May 2024 12:35:45 +0100 Subject: [PATCH] Addressed issue with disabled features in certain regions (#1618) ## Changes - Added error handling when listing serving endpoints, as it will raise NotFound error if the feature is disabled. ### Linked issues Resolves #1617 ### Tests - [x] manually tested - [x] added unit tests --- src/databricks/labs/ucx/workspace_access/generic.py | 7 +++++-- tests/unit/workspace_access/test_generic.py | 10 ++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/src/databricks/labs/ucx/workspace_access/generic.py b/src/databricks/labs/ucx/workspace_access/generic.py index 9818c861f7..deae7e2b53 100644 --- a/src/databricks/labs/ucx/workspace_access/generic.py +++ b/src/databricks/labs/ucx/workspace_access/generic.py @@ -55,8 +55,11 @@ def object_types(self) -> set[str]: def __iter__(self): started = datetime.datetime.now() - for item in self._func(): - yield GenericPermissionsInfo(getattr(item, self._id_attribute), self._object_type) + try: + for item in self._func(): + yield GenericPermissionsInfo(getattr(item, self._id_attribute), self._object_type) + except NotFound as e: + logger.error(f"Listing {self._object_type} failed: {e}") since = datetime.datetime.now() - started logger.info(f"Listed {self._object_type} in {since}") diff --git a/tests/unit/workspace_access/test_generic.py b/tests/unit/workspace_access/test_generic.py index c369ff62a3..c3cde2733f 100644 --- a/tests/unit/workspace_access/test_generic.py +++ b/tests/unit/workspace_access/test_generic.py @@ -916,3 +916,13 @@ def test_models_page_listing(): for item in auth_items: assert item.object_id == "/root" assert item.object_type == "registered-models" + + +def test_serving_endpoints_not_enabled(caplog): + ws = create_autospec(WorkspaceClient) + ws.serving_endpoints.list.side_effect = NotFound("Model serving is not enabled for your shard") + + sup = GenericPermissionsSupport(ws=ws, listings=[Listing(ws.serving_endpoints.list, "id", "serving-endpoints")]) + with caplog.at_level('ERROR'): + list(sup.get_crawler_tasks()) + assert "Listing serving-endpoints failed: Model serving is not enabled for your shard" in caplog.text