チュートリアル: Databricks アプリを Lakebase 自動スケールに接続する

Important

Lakebase 自動スケールは、ベータで、eastus2、westeurope、westusの各リージョンで利用可能です。

Lakebase 自動スケーリングは、自動スケール コンピューティング、ゼロへのスケーリング、分岐、インスタント リストアを備えた最新バージョンの Lakebase です。 Lakebase Provisioned との機能の比較については、バージョンの選択を参照してください。

このチュートリアルでは、Databricks アプリを、資格情報の自動ローテーションを使用して Lakebase 自動スケールに接続する方法について説明します。 アプリは、有効期限が切れる前に Databricks から新しいデータベース資格情報を生成します。 この例では Flask を使用していますが、認証パターンは任意のフレームワークに適用されます。

動作方法

Databricks Apps は、1 時間後に期限切れになる OAuth トークンを使用して Lakebase に対して認証を行います。 これを処理するには、アプリの サービス プリンシパル用に Postgres ロールを作成し、データベースに接続する必要がある場合は常に新しいトークンを自動的に生成するようにアプリを構成します。 これは、接続プール パターンによって発生します。プールは、必要に応じて新しいトークンを使用して新しい接続を作成するため、アプリで有効期限切れの資格情報が使用されることはありません。

Databricks にアプリをデプロイすると、アプリはサービス プリンシパルとして実行され、その ID のトークンが生成されます。 ローカルでテストすると、アプリは Databricks ユーザー アカウントとして実行され、トークンが生成されます。 どちらも同じトークン ローテーション コードを使用します。認証コンテキストのみが変更されます。

開始する前に

このチュートリアルを完了するには、次のものが必要です。

• Lakebase Postgres 自動スケーリングが有効になっている Databricks ワークスペースへのアクセス (アプリ スイッチャーに Lakebase が表示されない場合は、ワークスペース管理者に問い合わせてください)
• Databricks Apps を作成するためのアクセス許可
• Python と SQL に関する基本的な知識
• Databricks CLI がインストールされた ローカル開発用
• Python 3.9 以降がローカルにインストールされている

手順 1: アプリとデータベースを作成する

まず、Databricks アプリと Lakebase プロジェクトの両方を作成します。 アプリは、データベース認証に使用するサービス プリンシパル ID を自動的に取得します。

アプリを作成する

Flask Hello World テンプレート (+ ワークスペース内の新しい > アプリ ) を使用して、新しい Databricks アプリを作成します。 詳細な手順については、「 Databricks アプリの作成」を参照してください。

インストール後、アプリの [ 環境 ] タブに移動し、 DATABRICKS_CLIENT_ID 値 ( 6b215d2b-f099-4bdb-900a-60837201ecec などの UUID 形式) を書き留めます。 これは、OAuth 認証用のアプリの Postgres ユーザー名になります。

注

アプリをまだデプロイしないでください。最初にデータベース接続を構成します。

データベースを作成する

データベースをホストする新しい Lakebase 自動スケール プロジェクトを作成します。 [右上隅にあるアプリで、Lakebase を選択し、目的の名前 (my-app-db など) と Postgres バージョン (既定の Postgres 17 をそのまま使用) で新しいプロジェクトを作成します。 セットアップの詳細については、「 プロジェクトの作成」を参照してください。

コンピューティングがアクティブになるまで (約 1 分) 待ってから続行します。

手順 2: データベース認証とスキーマを構成する

OAuth 認証を使用してアプリのサービス プリンシパルの Postgres ロールを作成し、アプリに表示されるデータを含むサンプル テーブルを作成します。

OAuth 認証を設定する

Lakebase プロジェクトで、SQL エディターを開き、次のコマンドを実行します。 databricks_auth拡張機能では OAuth 認証が有効になり、Postgres ロールは従来のパスワードではなく Databricks トークンを受け入れることが可能になります。

-- Enable the Databricks authentication extension
CREATE EXTENSION IF NOT EXISTS databricks_auth;

-- Create a Postgres role for your app's service principal
-- Replace the UUID below with your DATABRICKS_CLIENT_ID from Step 1
SELECT databricks_create_role('<DATABRICKS_CLIENT_ID>', 'service_principal');

-- Grant necessary permissions (use the same DATABRICKS_CLIENT_ID)
GRANT CONNECT ON DATABASE databricks_postgres TO "<DATABRICKS_CLIENT_ID>";
GRANT CREATE, USAGE ON SCHEMA public TO "<DATABRICKS_CLIENT_ID>";

<DATABRICKS_CLIENT_ID>をアプリのDATABRICKS_CLIENT_ID値に置き換えます。 サービス プリンシパルは、Databricks が自動的に管理する OAuth トークンを使用して認証できるようになりました。 詳細については、「 Databricks ID の OAuth ロールを作成する」を参照してください。

データベース スキーマを作成する

サービス プリンシパルの明示的なアクセス許可を持つサンプル テーブルを作成します (サービス プリンシパルは既定のスキーマアクセス許可を継承しません)。

-- Create a sample table
CREATE TABLE notes (
    id SERIAL PRIMARY KEY,
    content TEXT NOT NULL,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);

-- Grant permissions to your app's service principal (use your DATABRICKS_CLIENT_ID)
GRANT SELECT, INSERT, UPDATE, DELETE ON TABLE notes TO "<DATABRICKS_CLIENT_ID>";

-- Insert sample data
INSERT INTO notes (content) VALUES
   ('Welcome to Lakebase Autoscaling!'),
   ('This app connects to Postgres'),
   ('Data fetched from your database');

<DATABRICKS_CLIENT_ID>をDATABRICKS_CLIENT_ID値に置き換えます。

手順 3: アプリケーションをビルドして構成する

アプリ ファイルをダウンロードし、OAuth トークンの自動ローテーションを使用してデータベース接続を構成し、デプロイする前にローカルでテストします。

アプリ ファイルをダウンロードして構成する

アプリの [ファイルの同期] セクションからエクスポート コマンドをコピーして、ワークスペースからアプリ ファイルを ダウンロードします。

databricks workspace export-dir /Workspace/Users/<your-email>/databricks_apps/<app-folder>/flask-hello-world-app .

app.yamlを編集して、データベース接続の詳細を追加します。 [パラメーターのみ] を選択して、Lakebase Connect モーダルから接続値を取得します。

command: ['flask', '--app', 'app.py', 'run', '--host', '0.0.0.0', '--port', '8000']

env:
  - name: PGHOST
    value: '<your-endpoint-hostname>'
  - name: PGDATABASE
    value: 'databricks_postgres'
  - name: PGUSER
    value: '<DATABRICKS_CLIENT_ID>'
  - name: PGPORT
    value: '5432'
  - name: PGSSLMODE
    value: 'require'
  - name: ENDPOINT_NAME
    value: 'projects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id>'

プレースホルダーを置き換えてください。

• <your-endpoint-hostname>: Connect モーダルから PGHOST 値をコピーします (例: ep-xyz.database.us-west-2.dev.databricks.com)
• <DATABRICKS_CLIENT_ID>: 手順 1 の DATABRICKS_CLIENT_ID を使用する
• projects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id>: databricks postgres list-endpoints を実行し、 name フィールドの値をコピーします (形式: projects/<uuid>/branches/<uuid>/endpoints/<id>)

OAuth トークンのローテーションとデータベース クエリを実装する

app.pyをこのコードに置き換えます。 これにより、認証用の OAuth トークンの自動ローテーションと、手順 2 で作成したメモをフェッチして表示するデータベース クエリという 2 つの重要な変更が実装されます。

import os
from databricks.sdk import WorkspaceClient
import psycopg
from psycopg_pool import ConnectionPool
from flask import Flask

app = Flask(__name__)

# Initialize Databricks client for token generation
w = WorkspaceClient()

# Custom connection class that generates fresh OAuth tokens
class OAuthConnection(psycopg.Connection):
    @classmethod
    def connect(cls, conninfo='', **kwargs):
        # Generate a fresh OAuth token for each connection (tokens are workspace-scoped)
        endpoint_name = os.environ["ENDPOINT_NAME"]
        credential = w.postgres.generate_database_credential(endpoint=endpoint_name)
        kwargs['password'] = credential.token
        return super().connect(conninfo, **kwargs)

# Configure connection parameters
username = os.environ["PGUSER"]
host = os.environ["PGHOST"]
port = os.environ.get("PGPORT", "5432")
database = os.environ["PGDATABASE"]
sslmode = os.environ.get("PGSSLMODE", "require")

# Create connection pool with automatic token rotation
pool = ConnectionPool(
    conninfo=f"dbname={database} user={username} host={host} port={port} sslmode={sslmode}",
    connection_class=OAuthConnection,
    min_size=1,
    max_size=10,
    open=True
)

@app.route('/')
def hello_world():
    # Use connection from pool (automatically gets fresh token)
    with pool.connection() as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT content, created_at FROM notes ORDER BY created_at DESC LIMIT 5")
            notes = cur.fetchall()

    # Display results
    notes_html = "<ul>" + "".join([f"<li>{note[0]} - {note[1]}</li>" for note in notes]) + "</ul>"
    return f'<h1>Hello from Lakebase!</h1><h2>Recent Notes:</h2>{notes_html}'

if __name__ == '__main__':
    app.run(host="0.0.0.0", port=8000)

このコードでは、次の 3 つの主要コンポーネントを使用して、データベース資格情報の自動ローテーションを実装します。

• WorkspaceClient: Databricks SDK を使用して新しいデータベース資格情報を生成します
• OAuthConnection クラス: 新しい接続ごとに新しい資格情報を挿入するカスタム接続クラス
• ConnectionPool: 接続の作成時にデータベース接続を管理し、カスタム接続クラスを呼び出します

このパターンにより、アプリは常に有効なデータベース資格情報を持ちます。 資格情報のローテーションのしくみ、さまざまなローテーション戦略、およびエラー処理の詳細については、「 Lakebase でのトークンのローテーション」を参照してください。

必要なパッケージを含むように requirements.txt を更新します。

flask
psycopg[binary,pool]
databricks-sdk>=0.81.0

最小 SDK バージョン (0.81.0) を使用すると、OAuth トークンの生成に generate_database_credential() メソッドを使用できるようになります。

ローカルでテストする

アプリをローカルでテストし、デプロイする前にデータベース接続が機能することを確認します。 ローカルでテストする場合、アプリは (サービス プリンシパルではなく) Databricks ユーザー アカウントとして実行されるため、次の環境変数で PGUSER をメール アドレスに変更する必要があります。

ワークスペースに対して認証を行い、環境変数をエクスポートします。

databricks auth login

export PGHOST="<your-endpoint-hostname>"
export PGDATABASE="databricks_postgres"
export PGUSER="your.email@company.com"  # Use YOUR email for local testing, not the service principal
export PGPORT="5432"
export PGSSLMODE="require"
export ENDPOINT_NAME="<your-endpoint-name>"

app.yamlから値をコピーしますが、PGUSER値 (サービス プリンシパル クライアント ID) を Databricks のメール アドレスに置き換えます。

依存関係をインストールし、アプリを実行します。

pip3 install --upgrade -r requirements.txt
python3 app.py

ブラウザーで http://localhost:8000 を開きます。 3 つのサンプル ノートと共に "Hello from Lakebase!" が表示されます。 接続プールは、新しい接続の作成時に新しい OAuth トークンを自動的に生成します。 詳細については、「 Lakebase での認証」を参照してください。

手順 4: デプロイして確認する

ローカルでテストした後、ワークスペース フォルダーに変更を同期し、その場所からデプロイします。

# Upload files to workspace
databricks sync . /Workspace/Users/<your-email>/my-lakebase-app

# Deploy from the uploaded location
databricks apps deploy <app-name> --source-code-path /Workspace/Users/<your-email>/my-lakebase-app

<your-email>を Databricks のメール アドレスに置き換え、<app-name>をアプリ名に置き換えます。 --source-code-path フラグは、アプリの既定の場所ではなく、アップロードしたファイルを使用するようにデプロイに指示します。

デプロイが完了するまで (2 ~ 3 分) 待ってから、指定された URL でアプリにアクセスします。 サンプル ノートと共に "Hello from Lakebase!" が表示されます。

こちらも参照ください

• Lakebase での認証
• Databricks Apps のドキュメント