Conexión de una aplicación personalizada de Databricks a Lakebase

En este tutorial se muestra cómo conectar una aplicación de Databricks al escalado automático de Lakebase con rotación automática de credenciales. La aplicación genera credenciales de base de datos nuevas de Databricks antes de que expiren. En el ejemplo se usa Flask, pero el patrón de autenticación se aplica a cualquier marco de trabajo.

Cómo funciona

Databricks Apps se autentica en Lakebase mediante tokens de OAuth que expiran después de una hora. Para controlar esto, crea un rol de Postgres para la entidad de servicio de la aplicación y, a continuación, configura la aplicación para generar automáticamente tokens nuevos siempre que necesite conectarse a la base de datos. La aplicación usa un grupo de conexiones. El grupo crea nuevas conexiones con tokens nuevos según sea necesario, por lo que la aplicación nunca usa credenciales expiradas.

Al implementar la aplicación en Azure Databricks, se ejecuta como su principal de servicio y genera tokens para esa identidad. Al probar localmente, la aplicación se ejecuta como la cuenta de usuario de Azure Databricks y genera tokens automáticamente. Ambos usan el mismo código de rotación de tokens. Solo cambia el contexto de autenticación.

Antes de empezar

Para completar este tutorial, necesita lo siguiente:

• Acceso a un área de trabajo de Azure Databricks con el escalado automático de Lakebase Postgres habilitado. Póngase en contacto con el administrador del área de trabajo si no ve Lakebase en el conmutador de la aplicación.
• Permiso para crear aplicaciones
• Conocimientos básicos de Python y SQL
• CLI de Databricks instalada para el desarrollo local
• Python 3.9 o posterior instalado localmente

Paso 1: Crear la aplicación y la base de datos

En primer lugar, cree una aplicación de Databricks y un proyecto de Lakebase. La aplicación obtiene automáticamente una identidad de principal de servicio que utilizará para la autenticación de la base de datos.

Creación de la aplicación

Cree una nueva aplicación de Databricks mediante la plantilla Hello world de Flask . Consulte Creación de una aplicación de Databricks a partir de una plantilla.

Después de crear la aplicación, vaya a la pestaña Entorno de la aplicación y anote el DATABRICKS_CLIENT_ID valor (formato UUID como 6b215d2b-f099-4bdb-900a-60837201ecec). Esto se convierte en el nombre de usuario de Postgres de la aplicación para la autenticación de OAuth.

Creación de la base de datos

Cree un nuevo proyecto de escalado automático de Lakebase para hospedar la base de datos. Haga clic en el Conmutador de aplicaciones, seleccione Lakebase Postgres y, a continuación, cree un nuevo proyecto con un nombre (por ejemplo, my-app-db) y la versión de Postgres (acepte el postgres 17 predeterminado). Para obtener información completa sobre la configuración, consulte Creación de un proyecto.

Espere a que el proceso se active (aproximadamente 1 minuto) antes de continuar.

Paso 2: Configuración de la autenticación y el esquema de la base de datos

Cree un rol de Postgres para la entidad de servicio de su aplicación con autenticación OAuth y, a continuación, cree una tabla de ejemplo con datos para que la aplicación los muestre.

Configuración de la autenticación de OAuth

En el proyecto de Lakebase, abra el Editor de SQL y ejecute estos comandos. La databricks_auth extensión habilita la autenticación de OAuth. Con él, los roles de Postgres aceptan tokens de Databricks en lugar de contraseñas tradicionales:

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

Reemplace <DATABRICKS_CLIENT_ID> por el valor de su DATABRICKS_CLIENT_ID aplicación. La entidad de servicio ahora puede autenticarse mediante tokens de OAuth que Azure Databricks administra automáticamente. Para obtener más información, consulte Crear un rol de OAuth para una identidad de Azure Databricks.

Creación de un esquema de base de datos

Cree una tabla de ejemplo con permisos explícitos para la entidad de servicio (las entidades de servicio no heredan los permisos de esquema predeterminados):

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

Reemplace <DATABRICKS_CLIENT_ID> con su DATABRICKS_CLIENT_ID valor.

Paso 3: Compilar y configurar la aplicación

Descargue los archivos de la aplicación, configure la conexión de base de datos con la rotación automática de tokens de OAuth y pruebe localmente antes de la implementación.

Descarga y configuración de archivos de aplicación

Descargue los archivos de la aplicación desde el área de trabajo copiando el comando de exportación desde la sección Sincronizar los archivos de la aplicación:

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

Edite app.yaml para agregar detalles de conexión de base de datos. Para obtener los valores de conexión del modal de Lakebase Connect , seleccione Solo parámetros:

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

Reemplace los marcadores de posición:

• <your-endpoint-hostname>: copie el valor PGHOST del modal Connect (por ejemplo, ep-xyz.database.us-west-2.dev.databricks.com)
• <DATABRICKS_CLIENT_ID>: Usa tu DATABRICKS_CLIENT_ID del paso 1.
• : En la aplicación Lakebase, vaya a la pestaña Computación de su sucursal, haga clic en Obtener IDCopiar nombre del recurso.

Implementación de la rotación de tokens de OAuth y la consulta de base de datos

Reemplace app.py por este código, que agrega la rotación automática de tokens de OAuth y una consulta de base de datos que captura las notas del paso 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)

Los tres componentes clave son:

• WorkspaceClient: genera credenciales nuevas mediante el SDK.
• OAuthConnection: una clase de conexión personalizada que inserta una credencial nueva en cada conexión.
• ConnectionPool: administra las conexiones y llama a la clase personalizada según sea necesario.

Para obtener más información sobre las estrategias de rotación de credenciales y el control de errores, consulte Ejemplos de rotación de tokens.

Actualice requirements.txt para incluir los paquetes necesarios:

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

La versión 0.81.0 o posterior incluye el generate_database_credential() método .

Probar localmente

Pruebe la aplicación localmente para comprobar que la conexión de base de datos funciona antes de la implementación. Al probar localmente, la aplicación se ejecuta como la cuenta de usuario de Azure Databricks (no la entidad de servicio), por lo que reemplaza PGUSER con tu dirección de correo electrónico en las variables de entorno siguientes.

Autentíquese en el área de trabajo y exporte variables de entorno:

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

Copie los valores de su app.yaml, pero reemplace el valor de PGUSER (identificador de cliente del principal de servicio) con su dirección de correo electrónico de Azure Databricks.

Instale las dependencias y ejecute la aplicación:

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

Abra http://localhost:8000 en el explorador. Debería ver "¡Hola desde Lakebase!" con sus tres notas de ejemplo. El grupo de conexiones genera automáticamente tokens de OAuth nuevos al crear nuevas conexiones. Para más información, consulte Autenticación de tokens de OAuth.

Paso 4: Implementar y comprobar

Después de probar localmente, sincronice los cambios en una carpeta del área de trabajo e implemente desde esa ubicació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

Reemplace <your-email> por su dirección de correo electrónico de Azure Databricks y <app-name> por el nombre de la aplicación. La --source-code-path marca indica a la implementación que use los archivos cargados en lugar de la ubicación predeterminada de la aplicación.

Espere a que la implementación se complete (de 2 a 3 minutos) y, a continuación, acceda a la aplicación en la dirección URL proporcionada. Deberías ver "¡Hola desde Lakebase!" con las notas de ejemplo.

Consulte también

• Ejemplos de rotación de tokens
• Aplicaciones de Databricks