Connectez une application Databricks personnalisée à Lakebase

Ce tutoriel vous montre comment connecter une application Databricks à Lakebase Autoscaling avec rotation automatique des identifiants. L'application génère de nouveaux identifiants de base de données à partir de Databricks avant leur expiration. L'exemple utilise Flask , mais le modèle d'authentification s'applique à tout framework.

Comment cela fonctionne

Les Databricks Apps s'authentifient auprès de Lakebase à l'aide de jetons OAuth qui expirent après une heure. Pour ce faire, vous créez un rôle Postgres pour le Service Principal de votre application, puis vous configurez votre application pour qu'elle génère automatiquement de nouveaux jetons chaque fois qu'elle doit se connecter à la base de données. L'application utilise un Pool de connexions. Le Pool crée de nouvelles connexions avec des jetons frais au besoin, de sorte que votre application n'utilise jamais de références expirées.

Lorsque vous déployez l'application sur Databricks, elle s'exécute en tant que son Service Principal et génère des jetons pour cette identité. Lorsque vous testez localement, l'application s'exécute en tant que votre compte utilisateur Databricks et génère des jetons pour vous. Les deux utilisent le même code de rotation de jeton. Seul le contexte d'authentification change.

Avant de commencer

Pour terminer ce tutoriel, vous avez besoin :

• Accès à un Workspace Databricks avec Lakebase Postgres Autoscaling activé. Contactez l'administrateur de votre Workspace si vous ne voyez pas Lakebase dans le sélecteur d'applications.
• Autorisation de créer des applications
• Familiarité de base avec Python et SQL
• Databricks CLI installé pour le développement local
• Python 3.9 ou version ultérieure installée localement

Étape 1 : Créez votre application et votre base de données.

Tout d'abord, créez une application Databricks et un projet Lakebase. L'application obtient automatiquement une identité de Service Principal que vous utiliserez pour l'authentification de base de données.

Créer l'application

Créez une nouvelle application Databricks à l'aide du Template Flask Hello world . Consultez Créer une application Databricks à partir d'un Template.

Après avoir créé l'application, accédez à l'onglet Environnement de l'application et notez la valeur DATABRICKS_CLIENT_ID (format UUID comme 6b215d2b-f099-4bdb-900a-60837201ecec). Ce nom d'utilisateur Postgres devient celui de votre application pour l'authentification OAuth.

remarque

Ne déployez pas encore l'application. Configurez d'abord la connexion à la base de données.

Créer la base de données

Créer un nouveau projet Lakebase de dimensionnement automatique pour héberger votre base de données. Cliquez sur le sélecteur d'applications, sélectionnez **Lakebase Postgres**, puis créez un nouveau projet avec un nom (par exemple,)my-app-db et une version de Postgres (acceptez le {default} Postgres 17). Pour tous les détails de configuration, consultez Créer un projet.

Attendez que le compute devienne actif (environ 1 minute) avant de poursuivre.

Étape 2 : Configurer l'authentification et le schéma de la base de données

Créez un rôle Postgres pour le Service Principal de votre application avec l'authentification OAuth, puis créez une table d'échantillon avec des données que votre application affichera.

Configurer l'authentification OAuth

Dans votre projet Lakebase, ouvrez l'éditeur SQL et exécutez ces commandes. L'extension databricks_auth active l'authentification OAuth. Ainsi, vos rôles Postgres acceptent les jetons Databricks au lieu des mots de passe traditionnels :

SQL

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

Remplacez <DATABRICKS_CLIENT_ID> par la valeur DATABRICKS_CLIENT_ID de votre application. Le Service Principal peut désormais s'authentifier à l'aide de jetons OAuth que Databricks gère automatiquement. Pour plus de détails, consultez Créer un rôle OAuth pour une identité Databricks.

Créer un schéma de base de données

Créez une table d'exemple avec des autorisations explicites pour votre Service Principal (les Service Principals n'héritent pas des autorisations de schéma par default) :

SQL

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

Remplacez <DATABRICKS_CLIENT_ID> par votre valeur DATABRICKS_CLIENT_ID.

Étape 3 : Créer et configurer votre application

Download vos fichiers d'application, configurez la connexion de la base de données avec la rotation automatique des jetons OAuth, et testez localement avant le déploiement.

download et configurer les fichiers de l'application

Download vos fichiers d'application du Workspace en copiant la commande d'exportation de la section Synchroniser les fichiers de l'application :

Bash

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

Modifier app.yaml pour ajouter les détails de connexion à la base de données. Obtenez vos valeurs de connexion à partir du modal **Connect** de Lakebase en sélectionnant **Paramètres uniquement** :

YAML

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

Remplacez les espaces réservés :

• <your-endpoint-hostname>: Copiez la valeur PGHOST de la fenêtre modale de connexion (par exemple, ep-xyz.database.us-west-2.dev.databricks.com)
• <DATABRICKS_CLIENT_ID>: Utilisez votre DATABRICKS_CLIENT_ID de l'étape 1
• projects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id>: Dans l’application Lakebase, accédez au tab **Computes** de votre Branch, cliquez sur **Obtenir l’ID** pour votre compute, et sélectionnez **Copier le nom de la ressource**.

Implémenter la rotation des jetons OAuth et la query de la base de données

Remplacez app.py par ce code, qui ajoute la rotation automatique des jetons OAuth et une query de base de données qui récupère les notes de l'étape 2 :

Python

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)

Les trois composants clés sont :

• WorkspaceClient: génère de nouvelles informations d'identification à l'aide du SDK.
• OAuthConnection: une classe de connexion personnalisée qui injecte un nouvel identifiant dans chaque connexion.
• ConnectionPool: gère les connexions et appelle la classe personnalisée si nécessaire.

Pour plus d'informations sur les stratégies de rotation des informations d'identification et la gestion des erreurs, consultez Exemples de rotation de jetons.

Mettre à jour requirements.txt pour inclure les packages requis :

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

La version 0.81.0 ou ultérieure inclut la méthode generate_database_credential().

Testez localement

Testez votre application localement pour vérifier que la connexion à la base de données fonctionne avant le déploiement. Lors des tests locaux, l'application s'exécute en tant que votre compte utilisateur Databricks (pas en tant que service principal), modifiez donc PGUSER par votre adresse e-mail dans les variables d'environnement ci-dessous.

Authentifiez-vous à votre Workspace et exportez les variables d'environnement :

Bash

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

Copiez les valeurs de votre app.yaml, mais remplacez la valeur PGUSER (ID client du Service Principal) par votre adresse e-mail Databricks.

Installer les dépendances et exécuter l'application :

Bash

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

Ouvrez http://localhost:8000 dans votre navigateur. Vous devriez voir « Hello from Lakebase ! » avec vos trois notes d'exemple. Le pool de connexions génère automatiquement de nouveaux jetons OAuth lors de la création de nouvelles connexions. Pour plus de détails, consultez l'authentification par jeton OAuth.

Étape 4 : Déployer et vérifier

Après les tests locaux, synchronisez vos modifications avec un dossier de workspace et déployez-les à partir de cet emplacement :

Bash

# 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

Remplacez <your-email> par votre adresse e-mail Databricks et <app-name> par le nom de votre application. L'indicateur --source-code-path indique au déploiement d'utiliser vos fichiers importés au lieu de l'emplacement par default de l'application.

Veuillez attendre la fin du déploiement (2 à 3 minutes), puis accédez à votre application à l'URL fournie. Vous devriez voir « Hello from Lakebase! » avec vos exemples de notes.

Ressources supplémentaires

• Exemples de rotation de jeton
• Databricks Apps