30 lines
1015 B
Python
30 lines
1015 B
Python
from typing import Dict
|
|
|
|
from .oauth import TokenSource
|
|
from .service.provisioning import Workspace
|
|
|
|
|
|
def add_workspace_id_header(cfg: "Config", headers: Dict[str, str]):
|
|
if cfg.azure_workspace_resource_id:
|
|
headers["X-Databricks-Azure-Workspace-Resource-Id"] = cfg.azure_workspace_resource_id
|
|
|
|
|
|
def add_sp_management_token(token_source: "TokenSource", headers: Dict[str, str]):
|
|
mgmt_token = token_source.token()
|
|
headers["X-Databricks-Azure-SP-Management-Token"] = mgmt_token.access_token
|
|
|
|
|
|
def get_azure_resource_id(workspace: Workspace):
|
|
"""
|
|
Returns the Azure Resource ID for the given workspace, if it is an Azure workspace.
|
|
:param workspace:
|
|
:return:
|
|
"""
|
|
if workspace.azure_workspace_info is None:
|
|
return None
|
|
return (
|
|
f"/subscriptions/{workspace.azure_workspace_info.subscription_id}"
|
|
f"/resourceGroups/{workspace.azure_workspace_info.resource_group}"
|
|
f"/providers/Microsoft.Databricks/workspaces/{workspace.workspace_name}"
|
|
)
|