Script for granting a Databricks user read-only access in a target workspace, using LHO

In this document we outline a python script which uses the Lakehouse Optimizer API to browse through all the published Azure subscriptions, Databricks workspaces, jobs, clusters and DLT pipelines and grant a user permissions for each of these assets:
attaching to clusters
viewing jobs
viewing DLT pipelines
Viewing SQL Warehouses
The script will use the LHO Service Principal with the admin role in each target workspace, called Admin Service Principal in this document.
secretScopeName parameter should be retrieved from LHO, Settings → Provisioning & Permissions → select Workspace → Secret Scope Enabled → Edit → Secret Scope Name
tenantId and clientId are on the same page at the bottom, Service Principal section
workspaceHost on the same page above the Service Principal section
Script file:
# Databricks notebook source
dbutils.widgets.text("username", "", "User to be granted")
dbutils.widgets.text("secretScopeName", "", "Secret Scope Name")
dbutils.widgets.text("tenantId", "", "Tenant Id")
dbutils.widgets.text("clientId", "", "Client Id")
dbutils.widgets.text("workspacesHost", "adb-***.0.azuredatabricks.net", "Workspaces (CSV)")
dbutils.widgets.text("dbx_api_throttle_timeout_in_ms", "50", "Dbx API Throttle Timeout")
dbutils.widgets.text("logLevel", "WARN", "Log level")

# COMMAND ----------

import requests
import os
import logging
import time
import re

# set up logging
log_level = dbutils.widgets.get("logLevel")
logging.basicConfig(level=log_level,
                    format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

userInfo = dbutils.widgets.get("username")
workspaces_hosts = [workspace.strip() for workspace in dbutils.widgets.get("workspacesHost").split(",")]

# login azure and obtain token for dbx calls
headers = {
    'Content-Type': 'application/x-www-form-urlencoded'
}

azure_databricks_scope = '2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default'


secretScopeName = dbutils.widgets.get("secretScopeName")
client_secret = dbutils.secrets.get(secretScopeName, "bplm-service-principal-clientSecret").strip()

if not client_secret:
    raise Exception("Client Secret not found in the secret scope")

# SP used to grant credentials. It should be admin in dbx
dbx_admin_sp_data = {
    'client_id': dbutils.widgets.get('clientId'),
    'grant_type': 'client_credentials',
    'scope': azure_databricks_scope,
    'client_secret': client_secret
}

databricks_api_throttle_in_ms = int(dbutils.widgets.get("dbx_api_throttle_timeout_in_ms"))
timeout = databricks_api_throttle_in_ms / 1000

log_level_order = ["ERROR", "WARN", "INFO", "DEBUG"]


def print_log(lvl, message):
    if log_level_order.index(lvl) <= log_level_order.index(log_level):
        print(f"{lvl}: {message}")


def debug(msg: str):
    print_log("DEBUG", msg)


def info(msg: str):
    print_log("INFO", msg)    


def warning(msg: str):
    print_log("WARN", msg)


def error(msg: str):
    print_log("ERROR", msg)

def host_to_dbx_id(host: str):
    # Extract the Workspace ID from the Databricks Host URL using regular expressions
    try:
        match = re.search(r"adb-(\d+)\..*", host)
        if match:
            workspace_id = match.group(1)
        else:
            workspace_id = None
    except Exception:
        workspace_id = None
    return workspace_id


def host_to_dbx_workspace(workspaces: list[str]):
    dbxWorkspaces = list(
        map(lambda host: (
            DatabricksWorkspace(Subscription("unknown", "unknown"), host_to_dbx_id(host), host, host, True))
            , workspaces)
    )
    return dbxWorkspaces
# Token for ServicePrincipal with admin access ( used in order to grant permissions)
dbx_admin_token_response = requests.post(
    f"https://login.microsoftonline.com/{dbutils.widgets.get('tenantId')}/oauth2/v2.0/token",
    headers=headers, data=dbx_admin_sp_data)
# This token will be used in "Grant" api requests
dbx_admin_sp_api_token = dbx_admin_token_response.json().get('access_token')


class Subscription:
    def __init__(self, subscription_id: str, display_name: str):
        self.subscription_id = subscription_id
        self.display_name = display_name

    def __str__(self):
        return f"{self.display_name}"


class DatabricksWorkspace:
    def __init__(self, subscription: Subscription, workspace_id: str, display_name: str, workspace_host: str,
                 is_premium: bool):
        self.subscription_id = subscription.subscription_id
        self.subscription_name = subscription.display_name
        self.workspace_id = workspace_id
        self.display_name = display_name
        self.workspace_host = workspace_host
        self.is_premium = is_premium

    def __str__(self):
        return f"{self.display_name}"


def process_error_response(api_response):
    response_text = api_response.headers.get('x-databricks-reason-phrase', api_response.text)

    if response_text is not None and len(api_response.text) > 0:
        processed_response = response_text
    else:
        processed_response = api_response.reason

    return f"[{api_response.status_code} - {processed_response}]"


def identity_to_be_granted():
    return userInfo


def get_grant_payload(permission: str):
    identity = identity_to_be_granted()
    identity_key = 'user_name'
    return {'access_control_list': [
        {
            identity_key: identity,
            'permission_level': permission
        }
    ]}


def process_grant_response(api_response, workspace_display_name: str, grant_type: str, entity_type: str,
                           entity_name: str):
    if api_response.status_code != 200:
        error(
            f"[workspace={workspace_display_name}] Could not grant {grant_type} to {identity_to_be_granted()} for {entity_type} {entity_name} [{api_response.status_code}] - {process_error_response(api_response)}")
    else:
        debug(
            f"[workspace={workspace_display_name}] Successfully granted {grant_type} to {identity_to_be_granted()} for {entity_type} {entity_name}")


# utility code
def grant_cluster_permission(host, display_name, cluster):
    debug(
        f"[workspace={display_name}] Granting permission CAN_ATTACH_TO to {identity_to_be_granted()} for cluster {cluster.get('cluster_name')}")

    payload = get_grant_payload('CAN_ATTACH_TO')
    api_response = requests.patch(f"https://{host}/api/2.0/permissions/clusters/{cluster.get('cluster_id')}",
                                  headers={'Authorization': f'Bearer {dbx_admin_sp_api_token}'}, json=payload)

    process_grant_response(api_response, display_name, "CAN_ATTACH_TO", "cluster", cluster.get('cluster_name'))


def grant_job_permission(host, display_name, job):
    job_id = job.get('job_id')
    job_name = job.get('settings', {}).get('name', job_id)

    print(
        f"[workspace={display_name}] Granting permission CAN_VIEW to {identity_to_be_granted()} for job {job_name}")

    payload = get_grant_payload('CAN_VIEW')
    api_response = requests.patch(f"https://{host}/api/2.0/permissions/jobs/{job_id}",
                                  headers={'Authorization': f'Bearer {dbx_admin_sp_api_token}'}, json=payload)

    process_grant_response(api_response, display_name, "CAN_VIEW", "job", job_name)


def grant_pipeline_permission(host, display_name, pipeline):
    debug(
        f"[workspace={display_name}] Granting permission CAN_VIEW to {identity_to_be_granted()} for pipeline {pipeline.get('name')}")

    payload = get_grant_payload('CAN_VIEW')
    api_response = requests.patch(f"https://{host}/api/2.0/permissions/pipelines/{pipeline.get('pipeline_id')}",
                                  headers={'Authorization': f'Bearer {dbx_admin_sp_api_token}'}, json=payload)

    process_grant_response(api_response, display_name, "CAN_VIEW", "pipeline", pipeline.get('name'))


def grant_warehouse_permission(host, display_name, warehouse):
    debug(
        f"[workspace={display_name}] Granting permission CAN_USE to {identity_to_be_granted()} for warehouse {warehouse.get('name')}")

    payload = get_grant_payload('CAN_USE')
    api_response = requests.patch(f"https://{host}/api/2.0/permissions/sql/warehouses/{warehouse.get('id')}",
                                  headers={'Authorization': f'Bearer {dbx_admin_sp_api_token}'}, json=payload)

    process_grant_response(api_response, display_name, "CAN_USE", "warehouse", warehouse.get('name'))


def process_clusters_in_workspace(workspace: DatabricksWorkspace):
    cluster_api_response = requests.get(f"https://{workspace.workspace_host}/api/2.0/clusters/list",
                                        headers={'Authorization': f'Bearer {dbx_admin_sp_api_token}'})
    if cluster_api_response.status_code != 200:
        error(
            f"[subscription={workspace.subscription_id}][workspace={workspace.display_name}] - Cluster request failure {process_error_response(cluster_api_response)}")
    else:
        clusters: list = cluster_api_response.json().get('clusters', [])
        if len(clusters) == 0:
            warning(
                f"[subscription={workspace.subscription_id}][workspace={workspace.display_name}] - No clusters in workspace. {cluster_api_response.json()}")

        # TODO filter to retrieve only APC
        filtered_clusters = list(filter(lambda cl: cl.get('cluster_source') != 'JOB', clusters))
        for cluster in filtered_clusters:
            time.sleep(timeout)
            grant_cluster_permission(workspace.workspace_host, workspace.display_name, cluster)


def process_all_jobs_in_workspace(workspace: DatabricksWorkspace):
    jobs = list()
    next_page_token = ""

    try:
        while next_page_token is not None:
            response = process_jobs_in_workspace(workspace, next_page_token)
            next_page_token = response.get("next_page_token", None)

            new_jobs = response.get("jobs", [])
            if len(new_jobs) == 0:
                warning(f"[subscription={workspace.subscription_id}][workspace={workspace.display_name}] - List jobs returned an empty list. {response}")
            jobs.extend(new_jobs)    
            time.sleep(timeout)
    except Exception as e:
        error(e)
    else:
        if len(jobs) == 0:
            warning(
                f"[subscription={workspace.subscription_id}][workspace={workspace.display_name}] - No jobs in workspace ")
        for job in jobs:
            time.sleep(timeout)
            grant_job_permission(workspace.workspace_host, workspace.display_name, job)


def process_jobs_in_workspace(workspace, next_page_token: str):
    if next_page_token is not None and next_page_token != "":
        token = f"&page_token={next_page_token}"
    else:
        token = ""

    jobs_api_response = requests.get(f"https://{workspace.workspace_host}/api/2.1/jobs/list?limit=25{token}",
                                     headers={'Authorization': f'Bearer {dbx_admin_sp_api_token}'})
    if jobs_api_response.status_code != 200:
        raise Exception(
            f"[subscription={workspace.subscription_id}][workspace={workspace.display_name}] - Jobs request failure {process_error_response(jobs_api_response)}")
    else:
        jobs_response: list = jobs_api_response.json()
        if len(jobs_response) == 0:
            warning(
                f"[subscription={workspace.subscription_id}][workspace={workspace.display_name}] - List jobs returned an empty list using offset {offset}. {jobs_response}")
        return jobs_response


def process_all_pipelines_in_workspace(workspace: DatabricksWorkspace):
    pipelines = list()
    next_page_token = ""

    try:
        while next_page_token is not None:
            pipelines_api_response = process_pipelines_in_workspace(workspace, next_page_token)
            next_page_token = pipelines_api_response.get('next_page_token', None)

            new_pipelines = pipelines_api_response.get('statuses', [])
            if len(new_pipelines) == 0:
                warning(
                    f"[subscription={workspace.subscription_id}][workspace={workspace.display_name}] - List pipelines returned an empty list. {pipelines_api_response}")

            pipelines.extend(new_pipelines)
            time.sleep(timeout)
    except Exception as e:
        error(e)
    else:
        if len(pipelines) == 0:
            warning(
                f"[subscription={workspace.subscription_id}][workspace={workspace.display_name}] - No pipelines in workspace")
        for pipeline in pipelines:
            time.sleep(timeout)
            grant_pipeline_permission(workspace.workspace_host, workspace.display_name, pipeline)


def process_pipelines_in_workspace(workspace: DatabricksWorkspace, next_page_token):
    if next_page_token is not None and next_page_token != "":
        token = f"&page_token={next_page_token}"
    else:
        token = ""

    pipelines_api_response = requests.get(
        f"https://{workspace.workspace_host}/api/2.0/pipelines?max_results=100{token}",
        headers={'Authorization': f'Bearer {dbx_admin_sp_api_token}'})
    if pipelines_api_response.status_code != 200:
        raise Exception(
            f"[subscription={workspace.subscription_id}][workspace={workspace.display_name}] - Pipelines request failure {process_error_response(pipelines_api_response)}")
    else:
        pipelines_response: dict = pipelines_api_response.json()
        return pipelines_response


def process_warehouses_in_workspace(workspace: DatabricksWorkspace):
    warehouses_api_response = requests.get(f"https://{workspace.workspace_host}/api/2.0/sql/warehouses",
                                           headers={'Authorization': f'Bearer {dbx_admin_sp_api_token}'})
    if warehouses_api_response.status_code != 200:
        error(
            f"[subscription={workspace.subscription_id}][workspace={workspace.display_name}] - Warehouses request failure {process_error_response(warehouses_api_response)}")
    else:
        warehouses: list = warehouses_api_response.json().get('warehouses', [])
        if len(warehouses) == 0:
            warning(
                f"[subscription={workspace.subscription_id}][workspace={workspace.display_name}] - No warehouses in workspace. {warehouses_api_response.json()}")
        for warehouse in warehouses:
            time.sleep(timeout)
            grant_warehouse_permission(workspace.workspace_host, workspace.display_name, warehouse)


def process_all_entities_in_workspace(workspace: DatabricksWorkspace):
    info(f"Processing all databricks entities in workspace {workspace}")
    #process_clusters_in_workspace(workspace)
    process_all_jobs_in_workspace(workspace)
    #process_all_pipelines_in_workspace(workspace)
    #process_warehouses_in_workspace(workspace)

    info(f"Finished processing all databricks entities in workspace {workspace}")

def main():
    info(f"Application configured to grant permissions for workspaces {workspaces_hosts}")
    workspaces = host_to_dbx_workspace(workspaces_hosts)
    for workspace in workspaces:
        process_all_entities_in_workspace(workspace)

    info("Processed all entities.\nProcess finished.")


if __name__ == '__main__':
    main()