Databrick REST APIによりDatabricksトークンを作成後にシークレットとして登録する方法

概要

Databrick REST APIによりDatabricksトークンを作成後にシークレットとして登録する方法を共有します。
Databricksノートブックにて下記の流れを自動で実施することで、トークンを安全に利用する方法を提示します。

1. Databricksトークンを手動で作成して変数にセット
2. REST APIによりDatabricksトークンを発行
3. DatabricksトークンをDatabricksシークレットとして登録
4. 不要なDatabricksトークンを削除

詳細は下記のGithub pagesのページをご確認ください。

コードを実行したい方は、下記のdbcファイルを取り込んでください。

https://github.com/manabian-/databricks_tecks_for_qiita/blob/main/tecks/create_sercrets_by_rest_api/create_sercrets_by_rest_api.dbc

実行環境

databricks runtime: 8.3.x-cpu-ml-scala2.12
Python version: 3.8.8
pyspark version: 3.1.2.dev0

手順

1. Databricksトークンを手動で作成して変数にセット

User Settingsにてtmpというコメントをつけたトークンを作成。

下記のコードのトークンの変数の値にトークンの値を設定。

import json

scope_name = "databricks_mlops" # scope名をセット
databricks_token = "databricks_token" # DatabricksトークンにおけるsecreatのKey名をセット
databricks_url = "databricks_url" # DatabricksのURLをsecretsのKey名をセット


# トークン
## 最初は下記にシークレットを格納
token = "dapi5c768bf24dbd3be15cbc728066cccccc"
# 2回目からはDatabricksシークレットから取得 
try:
      token = dbutils.secrets.get(scope_name, databricks_token)
except Exception:
    pass


# DatabricksのURL 
## 下記についてはJobsから実行する際には適切に値を取得できないことに注意
browserhostname = json.loads(dbutils.notebook.entry_point
                                      .getDbutils()
                                      .notebook()
                                      .getContext()
                                  .toJson()
                                 )['tags']['browserHostName']
db_url = f'https://{browserhostname}'

#3 2回目からはDatabricksシークレットから取得 
try:
    db_url = dbutils.secrets.get(scope_name, databricks_url)
except:
    pass

2. REST APIによりDatabricksトークンを発行

import requests

# 秒×分×時間×１月当たりの日数×月
lifetime_seconds = 60*60*24*30*3

# トークンに設定するコメント。削除対象の判断で利用。
comment = "by Rest API"

data = None

response = requests.post(
    f'{db_url}/api/2.0/token/create',
    headers={'Authorization': f'Bearer {token}'},
    json={
        "lifetime_seconds": lifetime_seconds,
        "comment": comment,
    },
)

if response.status_code == 200:
    created_token_value = response.json()['token_value']
    created_token_id = response.json()['token_info']['token_id']
elif response.status_code == 403:
    print(response.text)
else:
    print("Error geting the job: {0}: {1}".format(response.json()["error_code"],response.json()["message"]))

3. DatabricksトークンをDatabricksシークレットとして登録

## scopeを作成
import requests

# scope名をセット
scope_name = "databricks_mlops"

response = requests.post(
    f'{db_url}/api/2.0/secrets/scopes/create',
    headers={'Authorization': f'Bearer {token}'},
    json={
        "scope": scope_name,
    },
)

if response.status_code == 200:
    print(f'Success:create scope "{scope_name}"')
elif response.status_code == 400:
    # スコープが存在する場合の分岐
    if response.json().get('error_code') == "RESOURCE_ALREADY_EXISTS":
        print(f'scope "{scope_name}" already exists')
    else:
        print(response.text)
elif response.status_code == 403:
    print(response.text)
else:
    print("Error geting the job: {0}: {1}".format(response.json()["error_code"],response.json()["message"]))

# Databricksトークンをsecretsに登録
import requests

response = requests.post(
    f'{db_url}/api/2.0/secrets/put',
    headers={'Authorization': f'Bearer {token}'},
    json={
        "scope": scope_name,
        "key": "databricks_secrets",
        "string_value": created_token_value,
    },
)

if response.status_code == 200:
    print(f'Success:create or replace secrets {scope_name}')
elif response.status_code == 403:
    print(response.text)
else:
    print("Error geting the job: {0}: {1}".format(response.json()["error_code"],response.json()["message"]))

# DatabricksのURLをsecretsに登録
import requests

response = requests.post(
    f'{db_url}/api/2.0/secrets/put',
    headers={'Authorization': f'Bearer {token}'},
    json={
        "scope": scope_name,
        "key": "databricks_url",
        "string_value": db_url,
    },
)

if response.status_code == 200:
    print(f'Success:create or replace secrets {scope_name}')
elif response.status_code == 403:
    print(response.text)
else:
    print("Error geting the job: {0}: {1}".format(response.json()["error_code"],response.json()["message"]))

4. 不要なDatabricksトークンを削除

# 本処理作成したトークン以外のリストを作成。
# ただし、`tmp`と`by Rest API`というコメントが設定さているトークンのみ対象
import requests

response = requests.get(
    f'{db_url}/api/2.0/token/list',
    headers={'Authorization': f'Bearer {token}'},
)

if response.status_code == 200:
    input_dict = response.json()
    token_id_target_deletion = []
    for i,x in enumerate(input_dict['token_infos']):
        if   x['comment'] == 'tmp':
            token_id_target_deletion.append(x['token_id'])
        elif x['comment'] == comment and x['token_id'] != created_token_id:
            token_id_target_deletion.append(x['token_id'])

elif response.status_code == 403:
    print(response.text)
else:
    print("Error geting the job: {0}: {1}".format(response.json()["error_code"],response.json()["message"]))

# トークンを削除する関数を定義
import requests

def delte_databricks_token(token_id):
    response = requests.post(
        f'{db_url}/api/2.0/token/delete',
        headers={'Authorization': f'Bearer {token}'},
        json={
            "token_id": token_id,
        },
    )

    if response.status_code == 200:
        print(f'Success:delte token_id {token_id}')

    elif response.status_code == 403:
        print(response.text)
    else:
        print("Error geting the job: {0}: {1}".format(response.json()["error_code"],response.json()["message"]))

# 不要なトークンを削除
for token_id in token_id_target_deletion:
    delte_databricks_token(token_id)