Databricks SDKでも統合認証を使いたい

はじめに

Databricksでは統合認証と呼ばれる、ツールやSDKの言語に関わらず統一的な方法で認証を行う方法が用意されています。

その中でも、今回はなんだかんだ言って色々お世話になりそうな環境変数を使用した認証を試してみます。

ここで紹介されている環境変数はSDKだけでなく、CLIやTerraformでの認証にも使うことが出来ます。

認証を行うサンプルコード

今回はDatabricks SDK for Pythonを使用します。

from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

for c in w.catalogs.list():
    print(c.name)

みたいなコードを用意し、uvの環境変数読み込み機能を使用して環境変数経由で認証情報を注入します。

認証がうまくいけば、認証されたユーザ・サービスプリンシパルの権限に応じたカタログの一覧が見えるはずです。

以下のサンプルではDATABRICKS_AUTH_TYPEを使用して認証方法の指示をSDKに与えていますが、与えない場合は認証方法の優先度に応じて認証を試行します。

Databricks CLIの認証情報を使用した認証

Databricks CLIを使用してOAuth U2M認証を行ってから、その認証情報を使ってSDKの認証を行います。

使い方はとっても簡単で、databricks auth loginでCLI上でログインしてから以下の環境変数をセットするだけです。

DATABRICKS_CONFIG_PROFILEはDEFAULT以外のプロファイル名を使用したいときに設定します。

DATABRICKS_HOST=https://dbc-12345678-abcd.cloud.databricks.com
DATABRICKS_AUTH_TYPE=databricks-cli

Personal Access Tokenを使用した認証

PATを発行して、そのPATでアクセスします。

DATABRICKS_HOST=https://dbc-12345678-abcd.cloud.databricks.com
DATABRICKS_AUTH_TYPE=pat
DATABRICKS_TOKEN=<token>

サービスプリンシパルのシークレットを使用した認証

何らかの処理を自動化したい場合、サービスプリンシパルを使うことも多いと思います。

DATABRICKS_HOST=https://dbc-12345678-abcd.cloud.databricks.com
DATABRICKS_AUTH_TYPE=oauth-m2m
DATABRICKS_CLIENT_ID=00000000-0000-0000-0000-000000000000
DATABRICKS_CLIENT_SECRET=<secret>

おわり