【Databricks】GitLab CI/CDからWorkload Identity Federation (OIDC)で認証してDeclarative Automation Bundlesでジョブをデプロイしてみた

はじめに

データ事業本部のueharaです。

DatabricksへCI/CDからデプロイする際、Personal Access Token（PAT）やサービスプリンシパルのOAuthシークレットを利用する方法があります。

これらの認証情報を利用する場合、GitLab側でシークレットの保管と定期的なローテーションが必要です。

DatabricksのWorkload Identity Federationを利用すると、GitLab CI/CDが発行するOpenID Connect（OIDC）トークンをDatabricksの短期OAuthトークンと交換できます。

GitLab側にDatabricksのPATやOAuthシークレットを保存する必要はありません。

今回はGitLab CI/CDからWorkload Identity FederationでDatabricksへ認証し、Declarative Automation Bundles（旧Databricks Asset Bundles）でジョブをデプロイ＆実行してみたいと思います。

処理の流れ

今回の処理の流れは次のようになります。

GitLabにはDatabricksのAccount ID、Workspace URL、サービスプリンシパルのApplication IDを登録しますが、いずれもシークレットでは無いことがポイントになります。

ファイル準備

ディレクトリ構成

今回のディレクトリ構成は次の通りです。

.
├── .gitlab-ci.yml
├── databricks.yml
├── resources
│   └── sample_job.yml
└── src
    └── hello_notebook.py

databricks.yml でBundleとデプロイターゲットを定義し、 resources/sample_job.yml でデプロイするジョブを定義します。

.gitlab-ci.yml には、OIDCトークンの発行とBundleのデプロイ手順を記述します。

Bundleの設定

databricks.ymlには、Bundle名とデプロイターゲットを定義します。

bundle:
  name: gitlab_wif_poc

include:
  - resources/*.yml

targets:
  dev:
    mode: development
    default: true
  prod:
    mode: production

dev ターゲットには development モードを指定しています。

このモードでは、デプロイされるリソース名に実行主体（Databricksユーザー/サービスプリンシパル）を含むプレフィックスが自動的に追加されます。

prod ターゲットには production モードを指定しています。

サンプルジョブの定義

resources/sample_job.yml では、サーバーレスコンピュートでNotebookを実行するジョブを定義します。

resources:
  jobs:
    sample_job:
      name: ${bundle.name}_job
      tasks:
        - task_key: hello
          notebook_task:
            notebook_path: ../src/hello_notebook.py
            source: WORKSPACE
          environment_key: serverless_default
      environments:
        - environment_key: serverless_default
          spec:
            environment_version: "2"
            dependencies: []

environment_version には、サーバーレス環境のバージョンを指定します。

追加のPythonパッケージは利用しないため、今回 dependencies は空の配列としています。

Notebookの作成

ジョブから実行するsrc/hello_notebook.pyは次の通りです。

# Databricks notebook source
# MAGIC %md
# MAGIC # GitLab WIF PoC Notebook
# MAGIC GitLab CI/CD からワークロード ID フェデレーションでデプロイ・実行されるサンプル。

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

from datetime import datetime, timezone

print(f"Executed at: {datetime.now(timezone.utc).isoformat()}")

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

print(spark.sql("SELECT current_user() AS me").first()["me"])

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

df = spark.table("samples.nyctaxi.trips")
result = df.groupBy("pickup_zip").count().orderBy("count", ascending=False).limit(5)
display(result)

実行時刻（UTC）と current_user() を出力し、 samples カタログのデータを簡単に集計するスクリプトとなっています。

GitLab CI/CDの設定

.gitlab-ci.yml は次の通りです。

stages:
  - validate
  - deploy
  - run

variables:
  DATABRICKS_AUTH_TYPE: env-oidc

.databricks-job:
  image: ubuntu:24.04
  id_tokens:
    DATABRICKS_OIDC_TOKEN:
      aud: $DATABRICKS_ACCOUNT_ID
  before_script:
    - apt-get update -y
    - apt-get install -y curl ca-certificates unzip
    - curl -fsSL https://raw.githubusercontent.com/databricks/setup-cli/main/install.sh | sh
    - databricks version

validate_dev:
  extends: .databricks-job
  stage: validate
  rules:
    - if: $CI_COMMIT_BRANCH == "develop"
  script:
    - databricks current-user me
    - databricks bundle validate -t dev

validate_prod:
  extends: .databricks-job
  stage: validate
  rules:
    - if: $CI_COMMIT_BRANCH == "main"
  script:
    - databricks current-user me
    - databricks bundle validate -t prod

deploy_dev:
  extends: .databricks-job
  stage: deploy
  rules:
    - if: $CI_COMMIT_BRANCH == "develop"
  script:
    - databricks bundle deploy -t dev

run_dev:
  extends: .databricks-job
  stage: run
  rules:
    - if: $CI_COMMIT_BRANCH == "develop"
  script:
    - databricks bundle run sample_job -t dev

deploy_prod:
  extends: .databricks-job
  stage: deploy
  rules:
    - if: $CI_COMMIT_BRANCH == "main"
  script:
    - databricks bundle deploy -t prod

run_prod:
  extends: .databricks-job
  stage: run
  rules:
    - if: $CI_COMMIT_BRANCH == "main"
  script:
    - databricks bundle run sample_job -t prod

id_tokens でGitLabにOIDCトークンを発行させ、 DATABRICKS_OIDC_TOKEN へ格納します。

Audienceには、GitLabのCI/CDの変数から取得したDatabricksのAccount IDを指定します。

DATABRICKS_AUTH_TYPE に env-oidc を指定すると、Databricks CLIが DATABRICKS_OIDC_TOKEN を読み取り、DatabricksのOAuthトークンへ交換します。

develop と main で別々のvalidateジョブを定義し、それぞれ対応するBundleターゲットを検証します。

本番デプロイの条件は main を明示しています（※今回デフォルトブランチは developとする想定をしています）。

（Databricks)サービスプリンシパルの作成

ファイルの準備ができたので、次にGitLab CI/CDがDatabricksへアクセスするためのサービスプリンシパルを作成します。

Databricksのアカウントコンソールへアクセスし、左側のタブにある『User management』から『Service principals』を開きます。

『Add service principal』を選択し、任意の名前でサービスプリンシパルを作成します。

今回は gitlab-poc という名前にしました。

作成後に表示されるUUID形式のApplication IDを控えます。

※この値は、後ほどGitLabの DATABRICKS_CLIENT_ID へ設定します。

今回はWorkload Identity Federationを利用するため、OAuthシークレットは発行しません。

（Databricks）サービスプリンシパルをワークスペースに追加

作成したサービスプリンシパルを、デプロイ先のワークスペースへ追加します。

アカウントコンソールの『Workspaces』から対象ワークスペースを選択し、『Permissions』タブを開きます。

『Add permissions』を選択して、先ほど作成したサービスプリンシパルを追加します。

今回は検証用としてAdmin accessも付与をしています。

（Databricks）フェデレーションポリシーの作成

次に、GitLabが発行したOIDCトークンをサービスプリンシパルの認証に利用できるようにします。

アカウントコンソールの『User management』から作成したサービスプリンシパルを選択し、『Credentials & secrets』を開きます。

『Federation policies』タブの『Create policy』を選択します。

プロバイダーには『GitLab CI/CD』を選択します。

そうしたら、まずは main ブランチ用のポリシーを次の内容で作成します。

• Group：GitLabのGroup
• Project：GitLabのProject
• Ref type：Branch
• Branch name：main

※Group、Project、Ref type、Branch nameからSubjectが自動生成されます。

同じ手順で develop ブランチ用のポリシーも作成します。

• Group：（mainと同じ）
• Project：（mainと同じ）
• Ref type：Branch
• Branch name：develop

ここまで完了すると、以下のように2つのポリシーが作成されているかと思います。

これで main と develop から発行されたOIDCトークンを、作成したサービスプリンシパルとして認証できます。

Subjectは完全一致で照合されるため、別プロジェクトやポリシーに登録していないブランチのOIDCトークンは認証が通らない形になります。

（GitLab）プロジェクトの設定

Databricks側の設定ができたので、先に用意したディレクトリの内容をGitLabプロジェクトのリポジトリにpushを行います。

次に、GitLabの『Settings』 → 『CI/CD』 → 『プロジェクト変数』に進み、『変数の追加』を選択します。

変数には以下を登録します（いずれもシークレットではないのでMasked不要/Protectedは任意）。

なお、Protectedを有効にする場合、パイプラインを実行する develop および main ブランチは保護ブランチに設定されている必要があるのでご留意下さい。

ブランチが保護されていない場合、Protected変数はジョブへ渡されず、実行がエラーとなります。

パイプラインの実行・動作確認

準備が完了したため、 develop ブランチに適当なコミットを追加すると、validate_dev、deploy_dev、run_devの順にジョブが実行されます。

validate_dev では、出力されたユーザー名がサービスプリンシパルのApplication IDと一致しているため、GitLabのOIDCトークンを利用してサービスプリンシパルとして認証できていました。

続いて deploy_dev がBundleをデプロイし、 run_dev がデプロイしたジョブを実行します。

Databricksの『Jobs & Pipelines』を確認すると、developmentモードのプレフィックスが付いたジョブが作成されていました。

ジョブの実行結果を確認すると、Notebookが正常に完了しており、無事結果の確認もできました。

次に、 main ブランチでの動作も確認すべく、 develop → main のマージを行います。

パイプラインを確認すると、こちらも無事実行が成功しておりました。

再度Databricksの『Jobs & Pipelines』を確認すると、今度は想定通りプレフィックスが付いていないジョブが作成されています。

ジョブの実行結果を確認すると、こちらも無事結果の確認ができました。

最後に

GitLab CI/CDのOIDCトークンとDatabricksのWorkload Identity Federation を組み合わせ、PATやOAuthシークレットをGitLabへ保存せずにジョブをデプロイできました。

フェデレーションポリシーではGitLabのプロジェクトとブランチをSubjectとして制限できるため、許可するパイプラインをDatabricks側で指定できます。

長期間有効な認証情報の保管とローテーションが不要になるため、GitLab CI/CDからDatabricksへデプロイする場合の認証方法として利用しやすい構成かと思います。

参考になりましたら幸いです。

参考文献

• Enable workload identity federation for GitLab CI/CD
• Configure a federation policy
• Authentication for Declarative Automation Bundles
• GitLab: OpenID Connect（OIDC）Authentication Using ID Tokens