カスタム Databricks アプリを Lakebase に接続する

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

動作方法

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

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

始める前に

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

• Lakebase Postgres 自動スケーリングが有効になっているAzure Databricks ワークスペースへのアクセス。 アプリ スイッチャーに Lakebase が表示されない場合は、ワークスペース管理者に問い合わせてください。
• アプリを作成するためのアクセス許可
• 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 Postgres を選択し、名前 (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値に置き換えます。 サービス プリンシパルは、自動的に管理Azure Databricks OAuth トークンを使用して認証できるようになりました。 詳細については、「Azure 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>: Lakebase アプリで、ブランチの [コンピューティング ] タブに移動し、コンピューティングの [ID の取得 ] をクリックして、[ リソース名のコピー] を選択します。

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

app.pyを次のコードに置き換えます。これにより、OAuth トークンの自動ローテーションと、手順 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: SDK を使用して新しい資格情報を生成します。
• OAuthConnection: 各接続に新しい資格情報を挿入するカスタム接続クラス。
• ConnectionPool: 接続を管理し、必要に応じてカスタム クラスを呼び出します。

資格情報ローテーション戦略とエラー処理の詳細については、「 トークンローテーションの例」を参照してください。

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

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

バージョン 0.81.0 以降には、 generate_database_credential() メソッドが含まれています。

ローカルでテストする

アプリをローカルでテストし、デプロイする前にデータベース接続が機能することを確認します。 ローカルでテストする場合、アプリは (サービス プリンシパルではなく) Azure 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) をAzure Databricksの電子メール アドレスに置き換えます。

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

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

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

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

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

こちらも参照ください

• トークンローテーションの例
• Databricks Apps