From 4499c7a37085831dd17749a259a09375cc493795 Mon Sep 17 00:00:00 2001 From: Cor Zuurmond Date: Mon, 23 Sep 2024 13:18:33 +0200 Subject: [PATCH] Set `workspace_client` on `AccountContext` Should be solved properly by https://github.com/databrickslabs/ucx/issues/2507 --- src/databricks/labs/ucx/contexts/account_cli.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/databricks/labs/ucx/contexts/account_cli.py b/src/databricks/labs/ucx/contexts/account_cli.py index 20617edab3..55ff692be5 100644 --- a/src/databricks/labs/ucx/contexts/account_cli.py +++ b/src/databricks/labs/ucx/contexts/account_cli.py @@ -1,8 +1,7 @@ from functools import cached_property from os import environ -from databricks.sdk import AccountClient - +from databricks.sdk import AccountClient, WorkspaceClient from databricks.labs.ucx.account.aggregate import AccountAggregate from databricks.labs.ucx.account.metastores import AccountMetastores @@ -19,6 +18,14 @@ def __init__(self, ac: AccountClient, named_parameters: dict[str, str] | None = def account_client(self) -> AccountClient: return self._ac + @cached_property + def workspace_client(self) -> WorkspaceClient: + """Return any workspace client available""" + workspace_clients = self.account_workspaces.workspace_clients() + if len(workspace_clients) == 0: + return super().workspace_client + return workspace_clients[0] + @cached_property def workspace_ids(self): return [int(_.strip()) for _ in self.named_parameters.get("workspace_ids", "").split(",") if _]