From a211dfbc5bf24bef644b78dc07e1b7ae0189e01b Mon Sep 17 00:00:00 2001
From: Cor Zuurmond <jczuurmond@protonmail.com>
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 _]