Een aangepaste Databricks-app verbinden met Lakebase

In deze handleiding leert u hoe u een Databricks-app verbindt met Lakebase Autoscaling met automatische rotatie van credentials. De app genereert nieuwe databasereferenties van Databricks voordat ze verlopen. In het voorbeeld wordt Flask gebruikt, maar het verificatiepatroon is van toepassing op elk framework.

Hoe het werkt

Databricks-apps verifiëren bij Lakebase met behulp van OAuth-tokens die na één uur verlopen. Hiervoor maakt u een Postgres-rol voor de service-principal van uw app en configureert u vervolgens uw app om automatisch nieuwe tokens te genereren wanneer deze verbinding moet maken met de database. De app maakt gebruik van een verbindingsgroep. De pool maakt zo nodig nieuwe verbindingen met nieuwe tokens, zodat uw app nooit verlopen referenties gebruikt.

Wanneer u de app implementeert in Azure Databricks, draait deze als zijn eigen service-principal en genereert tokens voor die identiteit. Wanneer u lokaal test, wordt de app uitgevoerd als uw Azure Databricks gebruikersaccount en worden tokens voor u gegenereerd. Beide gebruiken dezelfde tokenrotatiecode. Alleen de verificatiecontext verandert.

Voordat u begint

U hebt het volgende nodig om deze zelfstudie te voltooien:

• Toegang tot een Azure Databricks werkruimte waarvoor Automatische schaalaanpassing van Lakebase Postgres is ingeschakeld. Neem contact op met uw werkruimtebeheerder als u Lakebase niet ziet in de app-switcher.
• Machtiging voor het maken van apps
• Basiskennis van Python en SQL
• Databricks CLI geïnstalleerd voor lokale ontwikkeling
• Python 3.9 of hoger lokaal geïnstalleerd

Stap 1: Uw app en database maken

Maak eerst een Databricks-app en een Lakebase-project. De app krijgt automatisch een service-principal-identiteit die u gaat gebruiken voor databaseverificatie.

De app maken

Maak een nieuwe Databricks-app met behulp van de Flask Hello world-sjabloon . Zie Een Databricks-app maken op basis van een sjabloon.

Nadat u de app hebt gemaakt, gaat u naar het tabblad Omgeving van de app en noteert u de DATABRICKS_CLIENT_ID waarde (UUID-indeling zoals 6b215d2b-f099-4bdb-900a-60837201ecec). Dit wordt de Postgres-gebruikersnaam van uw app voor OAuth-verificatie.

De database maken

Maak een nieuw Lakebase-project voor automatisch schalen om uw database te hosten. Klik op het App-switcher, selecteer Lakebase Postgres en maak vervolgens een nieuw project met een naam (bijvoorbeeld my-app-db) en Postgres-versie (accepteer de standaardPostgres 17). Zie Een project maken voor volledige installatiedetails.

Wacht tot de berekening actief is (ongeveer 1 minuut) voordat u doorgaat.

Stap 2: Configureer database-authenticatie en schema

Maak een Postgres-rol voor de service-principal van uw app met OAuth-verificatie en maak vervolgens een voorbeeldtabel met gegevens die uw app kan weergeven.

OAuth-verificatie instellen

Open de SQL-editor in uw Lakebase-project en voer deze opdrachten uit. Met de databricks_auth extensie wordt OAuth-verificatie ingeschakeld. Hiermee accepteren uw Postgres-rollen Databricks-tokens in plaats van traditionele wachtwoorden:

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

Vervang <DATABRICKS_CLIENT_ID> met de DATABRICKS_CLIENT_ID-waarde van uw app. nl-NL: De service-principal kan nu worden geverifieerd met behulp van OAuth-tokens die automatisch worden beheerd door Azure Databricks. Zie Een OAuth-rol maken voor een Azure Databricks identiteit voor meer informatie.

Databaseschema maken

Maak een voorbeeldtabel met expliciete machtigingen voor uw service-principal (service-principals nemen geen standaardschemamachtigingen over):

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

Vervang <DATABRICKS_CLIENT_ID> door uw DATABRICKS_CLIENT_ID waarde.

Stap 3: Uw toepassing bouwen en configureren

Download uw app-bestanden, configureer de databaseverbinding met automatische OAuth-tokenrotatie en test lokaal voordat u implementeert.

App-bestanden downloaden en configureren

Download uw app-bestanden uit de werkruimte door de exportopdracht te kopiëren uit de sectie Synchroniseren van de bestanden van de app:

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

Bewerken app.yaml om details van de databaseverbinding toe te voegen. Haal uw verbindingswaarden op uit de modale Lakebase Connect door alleen parameters te selecteren:

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

Vervang de plaatsaanduidingen:

• <your-endpoint-hostname>: Kopieer de PGHOST-waarde uit de modale verbinding (bijvoorbeeld ep-xyz.database.us-west-2.dev.databricks.com)
• <DATABRICKS_CLIENT_ID>: Gebruik uw DATABRICKS_CLIENT_ID gegevens uit stap 1
• projects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id>: Ga in de Lakebase-app naar het tabblad Computes van uw tak, klik op ID ophalen voor uw computing en selecteer Resourcenaam kopiëren.

OAuth-tokenrotatie en databasequery implementeren

Vervang app.py door deze code, waarmee automatische OAuth-tokenrotatie en een databasequery worden toegevoegd waarmee de notities uit stap 2 worden opgehaald:

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)

De drie belangrijkste onderdelen zijn:

• WorkspaceClient: genereert nieuwe referenties met behulp van de SDK.
• OAuthConnection: een aangepaste verbindingsklasse die een nieuwe referentie in elke verbinding injecteert.
• ConnectionPool: beheert verbindingen en roept indien nodig de aangepaste klasse aan.

Zie Voorbeelden van tokenrotatie voor meer informatie over strategieën voor referentierotatie en foutafhandeling.

Werk requirements.txt bij om de vereiste pakketten op te nemen:

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

Versie 0.81.0 of hoger bevat de generate_database_credential() methode.

Lokaal testen

Test uw app lokaal om te controleren of de databaseverbinding werkt voordat u implementeert. Bij het lokaal testen draait de app op uw Azure Databricks gebruikersaccount (niet de service-principal), dus vervang PGUSER door uw e-mailadres in de omgevingsvariabelen hieronder.

Verifiëren bij uw werkruimte en omgevingsvariabelen exporteren:

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

Kopieer de waarden uit uw app.yaml, maar vervang de PGUSER waarde (service-principal-client-id) door uw Azure Databricks e-mailadres.

Installeer afhankelijkheden en voer de app uit:

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

Open http://localhost:8000 in je browser. U ziet 'Hello from Lakebase!' met uw drie voorbeeldnotities. De verbindingsgroep genereert automatisch nieuwe OAuth-tokens bij het maken van nieuwe verbindingen. Zie OAuth-tokenverificatie voor meer informatie.

Stap 4: Implementeren en verifiëren

Nadat u lokaal hebt getest, synchroniseert u uw wijzigingen in een werkruimtemap en implementeert u deze vanaf die locatie:

# 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

Vervang <your-email> door uw Azure Databricks e-mailadres en <app-name> door de naam van uw app. De --source-code-path vlag vertelt de implementatie dat de geüploade bestanden moeten worden gebruikt in plaats van de standaardlocatie van de app.

Wacht tot de implementatie is voltooid (2-3 minuten) en open vervolgens uw app via de opgegeven URL. U ziet 'Hello from Lakebase!' met uw voorbeeldnotities.

Zie ook

• Voorbeelden van tokenrotatie
• Databricks-apps