Service PrincipalをDatabricksクラスタのシングルユーザに指定する

※　備忘録的な内容です。

仕事で必要になったため、Databricksでシングルノードかつセキュリティモード＝シングルユーザのAll Purposeクラスタを作り、かつシングルユーザはService Principalを指定しないといけない機会がありました。

ただ、DatabricksのUI上ではService Principalをユーザに指定できないんですよね。
ググったらコミュニティでも質問出てました。

CLIを使うのが一番簡単だとは思うのですが、REST APIを直接叩いても簡単に変更できます。今回はREST APIを叩いて変更する方法の記録です。

前提として、既にクラスタ自体は作成されており、セキュリティモードとシングルユーザ名だけを変えたいという状況とします。

Databricks上でノートブックを作成して、以下のコードをAPIのURLやcluster_id等を適切に設定して実行すると変更できます。

import requests

# 変更するクラスタのID
cluster_id = "xxxx-xxxxxx-xxxxxxx"
# 指定するService PrincipalのユーザID
single_user_name = "xxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxx"

# クラスタ設定を部分変更するAPIのURL
endpoint_change = f"https://xxxxxx.databricks.com/api/2.1/clusters/update"

# ノートブック実行ユーザからAPIトークンを取得
token = (
    dbutils.notebook.entry_point.getDbutils()
    .notebook()
    .getContext()
    .apiToken()
    .getOrElse(None)
)

# REST API実行のためのペイロードとヘッダ
headers = {"Authorization": "Bearer %s" % token, "Content-Type": "application/json"}
payload_change = {
    "cluster_id": cluster_id,
    "cluster": {
        "data_security_mode": "SINGLE_USER",
        "single_user_name": single_user_name,
    },
    "update_mask": "data_security_mode,single_user_name",
}

# 更新実行
response = requests.post(endpoint_change, headers=headers, json=payload_change)

# 結果の表示
response.json()

ちなみにCLIだとdatabricks clusters updateコマンドで同様の部分変更が可能だと思います。