教學：將 Databricks 應用程式連接到 Lakebase 自動縮放

這很重要

Lakebase 自動縮放在以下地區處於 Beta 階段：eastus2、westeurope、westus。

Lakebase 自動縮放是 Lakebase 的最新版本，具備自動縮放運算、縮放至零、分支及即時還原功能。關於與 Lakebase Provisioned 的功能比較，請參見「選擇版本」。

這個教學會教你如何將 Databricks 應用程式連接到 Lakebase 自動擴展，並自動進行憑證輪換。該應用程式會在資料庫憑證過期前從 Databricks 產生新的憑證。範例中使用Flask，但認證模式適用於任何框架。

運作方式

Databricks Apps 使用 OAuth 憑證向 Lakebase 進行認證，該憑證在一小時後過期。為了處理這個問題，你可以為應用程式的服務主體建立一個 Postgres 角色，然後設定應用程式在需要連接資料庫時自動產生新的標記。這是透過連線池模式實現的：連線池會根據需要建立新的連線並使用新的令牌，這樣你的應用程式就不會使用過期的憑證。

當你將應用程式部署到 Databricks 時，它將以其服務主體身份運行，並為該身份生成憑證。當你在本地測試時，應用程式會以你的 Databricks 使用者帳號執行，並為你產生代幣。兩者使用相同的令牌輪換代碼：只有認證上下文會改變。

開始之前

若要完成本教學課程，您需要：

啟用了 Lakebase Postgres 自動擴展功能的 Databricks 工作區（如果在應用程式切換器中看不到 Lakebase，請聯絡您的工作區管理員）
建立 Databricks 應用程式的權限
對 Python 和 SQL 的基本熟悉
為本地開發安裝的 Databricks CLI
本地安裝 Python 3.9 或更新版本

步驟一：建立您的應用程式與資料庫

首先，建立一個 Databricks 應用程式和一個 Lakebase 專案。應用程式會自動取得一個服務主體身份，你將用它來進行資料庫認證。

建立應用程式

使用 Flask Hello World 範本建立一個新的 Databricks 應用程式（+ > 在您的工作區新增應用程式）。詳細說明請參閱「建立 Databricks 應用程式」。

安裝後，到應用程式的環境標籤，記錄值（ DATABRICKS_CLIENT_ID UUID 格式如 6b215d2b-f099-4bdb-900a-60837201ecec）。這會成為你應用程式的 Postgres 用戶名，用於 OAuth 認證。

備註

先不要部署應用程式——你會先設定資料庫連線。

建立資料庫

建立一個新的 Lakebase 自動擴展專案來承載你的資料庫。點擊應用程式圖示。右上角的應用程式，選擇 Lakebase，然後建立一個新專案，名稱為你想要的名稱（例如 my-app-db）和 Postgres 版本（接受預設的 Postgres 17）。完整設定細節請參見「建立專案」。

等待運算啟動（約 1 分鐘）後再繼續。

步驟 2：設定資料庫認證與結構

為你的應用程式服務主體建立一個 Postgres 角色，並使用 OAuth 認證，然後建立一個範例表格，裡面有應用程式會顯示的資料。

設定 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 身份建立 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 步建立的筆記：

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)

本程式碼實作了自動資料庫憑證輪替，包含三個關鍵元件：

WorkspaceClient：透過 Databricks SDK 產生全新的資料庫憑證
OAuthConnection 類別：自訂連線類別，會為每個新連線注入新的憑證
ConnectionPool：管理資料庫連線並在建立連線時呼叫自訂連線類別

這種模式確保你的應用程式始終擁有有效的資料庫憑證。 關於憑證輪替的詳細說明、不同的輪替策略及錯誤處理，請參見 Lakebase 中的令牌輪替。

更新 requirements.txt 以包含所需套件：

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

最低 SDK 版本（0.81.0）確保該 generate_database_credential() 方法可用於 OAuth 代幣產生。

在本機進行測試

部署前先在本地測試你的應用程式，確認資料庫連線是否正常。在本地測試時，應用程式是以你的 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 。你應該會看到「Hello from Lakebase！」和你的三個樣本筆記。連線池在建立新連線時會自動產生新的 OAuth 代幣。更多細節請參閱 Lakebase 中的認證。

本地應用程式輸出顯示「來自 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

用您的 Databricks 電子郵件地址替換<your-email>，用您的應用程式名稱替換<app-name>。這個 --source-code-path 旗標會告訴部署使用你上傳的檔案，而不是應用程式的預設位置。

等待部署完成（2-3分鐘），然後以提供的網址存取應用程式。你應該會看到「Hello from Lakebase!」及你的樣本筆記。

另請參閱

意見反應

此頁面對您有幫助嗎？

Last updated on 2026-02-07