Připojení vlastní aplikace Databricks k Lakebase

V tomto kurzu se dozvíte, jak připojit aplikaci Databricks k automatickému škálování Lakebase pomocí automatické obměny přihlašovacích údajů. Aplikace vygeneruje z Databricks nové přihlašovací údaje databáze před vypršením jejich platnosti. Příklad používá Flask, ale vzor ověřování platí pro libovolnou architekturu.

Jak to funguje

Aplikace Databricks se ověřují ve službě Lakebase pomocí tokenů OAuth, jejichž platnost vyprší po jedné hodině. Abyste to zvládli, vytvoříte pro instanční objekt vaší aplikace roli Postgres a pak nakonfigurujete aplikaci tak, aby automaticky vygenerovala nové tokeny, kdykoli se potřebuje připojit k databázi. Aplikace používá fond připojení. Fond vytvoří nová připojení s čerstvými tokeny podle potřeby, takže vaše aplikace nikdy nepoužívá přihlašovací údaje s vypršenou platností.

Když nasadíte aplikaci do Azure Databricks, spustí se jako instanční objekt a vygeneruje tokeny pro tuto identitu. Když testujete místně, aplikace se spustí jako váš Azure Databricks uživatelský účet a vygeneruje tokeny za vás. Oba používají stejný kód obměně tokenů. Změní se pouze kontext ověřování.

Než začnete

K dokončení tohoto kurzu potřebujete:

Krok 1: Vytvoření aplikace a databáze

Nejprve vytvořte aplikaci Databricks i projekt Lakebase. Aplikace automaticky získá identitu hlavního objektu služby, kterou použijete pro ověření databáze.

Vytvoření aplikace

Vytvořte novou aplikaci Databricks pomocí šablony Flask Hello World . Viz Vytvoření aplikace Databricks ze šablony.

Po vytvoření aplikace přejděte na kartu Prostředí aplikace a poznamenejte si DATABRICKS_CLIENT_ID hodnotu (formát UUID, například 6b215d2b-f099-4bdb-900a-60837201ecec). To se stane uživatelským jménem Postgres vaší aplikace pro ověřování OAuth.

Poznámka:

Zatím aplikaci nenasazujte. Nejprve nakonfigurujte připojení k databázi.

Vytvoření databáze

Vytvořte nový projekt automatického škálování Lakebase pro hostování databáze. Klikněte na ikonu aplikace. Přepínač aplikací, vyberte Lakebase Postgres a pak vytvořte nový projekt s názvem (například my-app-db) a verzí Postgres (přijměte výchozí Postgres 17). Úplné podrobnosti o nastavení najdete v tématu Vytvoření projektu.

Než budete pokračovat, počkejte, než se výpočetní prostředky aktivují (přibližně 1 minutu).

Krok 2: Konfigurace ověřování a schématu databáze

Vytvořte roli Postgres pro instanční objekt vaší aplikace s ověřováním OAuth a pak vytvořte ukázkovou tabulku s daty, která se mají zobrazit v aplikaci.

Nastavení ověřování OAuth

V projektu Lakebase otevřete Editor SQL a spusťte tyto příkazy. Rozšíření databricks_auth umožňuje ověřování OAuth. Díky tomu vaše role Postgres přijímají tokeny Databricks místo tradičních hesel:

-- 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>";

Nahraďte <DATABRICKS_CLIENT_ID> hodnotou aplikace DATABRICKS_CLIENT_ID . Služební identita se teď může ověřovat pomocí tokenů OAuth, které Azure Databricks automaticky spravuje. Podrobnosti najdete v tématu Vytvoření role OAuth pro identitu Azure Databricks.

Vytvoření schématu databáze

Vytvořte ukázkovou tabulku s explicitními oprávněními pro instanční objekt (instanční objekty nedědí výchozí oprávnění schématu):

-- 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');

Nahraďte <DATABRICKS_CLIENT_ID> hodnotou DATABRICKS_CLIENT_ID .

Krok 3: Sestavení a konfigurace aplikace

Stáhněte si soubory aplikace, nakonfigurujte připojení k databázi pomocí automatické obměně tokenů OAuth a před nasazením otestujte místně.

Stažení a konfigurace souborů aplikací

Stáhněte si soubory aplikace z pracovního prostoru zkopírováním příkazu pro export z oddílu Synchronizovat soubory aplikace:

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

Upravte app.yaml, abyste přidali podrobnosti o připojení k databázi. Pokud chcete získat hodnoty připojení z modálního okna Lakebase Connect, vyberte Pouze parametry:

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>'

Vyměňte zástupné symboly:

  • <your-endpoint-hostname>: Zkopírujte hodnotu PGHOST z okna připojení (například ep-xyz.database.us-west-2.dev.databricks.com)
  • <DATABRICKS_CLIENT_ID>: Použijte DATABRICKS_CLIENT_ID z kroku 1
  • projects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id>: V aplikaci Lakebase přejděte na kartu Computes vaší větve, klikněte na Získat ID výpočetních prostředků a vyberte Kopírovat název prostředku.

Implementace obměně tokenů OAuth a databázového dotazu

Nahraďte app.py tímto kódem, který přidá automatické obměně tokenů OAuth a databázový dotaz, který načte poznámky z kroku 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)

Jedná se o tři klíčové komponenty:

  • WorkspaceClient: Generuje nové přihlašovací údaje pomocí sady SDK.
  • OAuthConnection: Vlastní třída připojení, která do každého připojení vloží nové přihlašovací údaje.
  • ConnectionPool: spravuje připojení a podle potřeby volá vlastní třídu.

Další informace o strategiích obměně přihlašovacích údajů a zpracování chyb najdete v příkladech obměně tokenů.

Aktualizujte requirements.txt , aby zahrnovaly požadované balíčky:

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

Verze 0.81.0 nebo novější obsahuje metodu generate_database_credential() .

Místní testování

Před nasazením otestujte aplikaci místně a ověřte, že připojení k databázi funguje. Při místním testování se aplikace spustí jako váš uživatelský účet Azure Databricks (ne jako služební principál), takže změňte PGUSER na svoji e-mailovou adresu v proměnných prostředí níže.

Ověřte se ve svém pracovním prostoru a exportujte proměnné prostředí:

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>"

Zkopírujte hodnoty z app.yaml, ale nahraďte hodnotu PGUSER (ID klienta instančního objektu) vaší Azure Databricks e-mailovou adresou.

Nainstalujte závislosti a spusťte aplikaci:

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

Otevřete http://localhost:8000 v prohlížeči. Měli byste vidět "Ahoj z Lakebase!" s vašimi třemi ukázkovými poznámkami. Fond připojení při vytváření nových připojení automaticky generuje nové tokeny OAuth. Další podrobnosti najdete v tématu Ověřování tokenu OAuth.

Výstup místní aplikace zobrazující

Krok 4: Nasazení a ověření

Po místním testování synchronizujte změny do složky pracovního prostoru a nasaďte je z daného umístění:

# 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

Nahraďte <your-email> e-mailovou adresou Azure Databricks a <app-name> názvem vaší aplikace. Příznak --source-code-path instruuje nasazení, aby používalo nahrané soubory místo výchozího umístění aplikace.

Počkejte, až se nasazení dokončí (2 až 3 minuty) a pak se k aplikaci dostanete na zadanou adresu URL. Měli byste vidět "Ahoj z Lakebase!" s ukázkovými poznámkami.

Další zdroje informací