Ligue uma aplicação Databricks personalizada ao Lakebase

Este tutorial mostra-lhe como ligar uma aplicação Databricks ao Lakebase Autoscaling com rotação automática de credenciais. A aplicação gera novas credenciais de base de dados a partir do Databricks antes de expirarem. O exemplo usa o Flask, mas o padrão de autenticação aplica-se a qualquer framework.

Como funciona

As aplicações Databricks autenticam-se no Lakebase usando tokens OAuth que expiram após uma hora. Para tratar disto, cria um papel Postgres para o principal de serviço da sua aplicação e depois configura a sua aplicação para gerar automaticamente novos tokens sempre que precisar de se ligar à base de dados. A aplicação utiliza um pool de ligações. O pool cria novas ligações com tokens novos conforme necessário, para que a sua aplicação nunca use credenciais expiradas.

Quando implementas a aplicação no Azure Databricks, ela corre como principal de serviço e gera tokens para essa identidade. Quando testas localmente, a aplicação corre como a tua conta de utilizador Azure Databricks e gera tokens para ti. Ambos usam o mesmo código de rotação de token. Apenas o contexto de autenticação muda.

Antes de começar

Para concluir este tutorial, você precisa:

• Acesso a um espaço de trabalho do Azure Databricks com Lakebase Postgres Autoscaling ativado. Contacta o administrador do teu espaço de trabalho se não vires o Lakebase no comutador de aplicações.
• Permissão para criar aplicações
• Familiaridade básica com Python e SQL
• Databricks CLI instalado para desenvolvimento local
• Python 3.9 ou posterior instalado localmente

Passo 1: Crie a sua aplicação e base de dados

Primeiro, cria tanto uma aplicação Databricks como um projeto Lakebase. A aplicação recebe automaticamente uma identidade de principal de serviço que irá usar para autenticação da base de dados.

Crie a aplicação

Crie uma nova aplicação Databricks usando o modelo Flask Helloworld . Consulte Criar um aplicativo Databricks a partir de um modelo.

Depois de criar a aplicação, vá ao separador Ambiente da aplicação e anote o DATABRICKS_CLIENT_ID valor (formato UUID como 6b215d2b-f099-4bdb-900a-60837201ecec). Isto torna-se o nome de utilizador Postgres da sua aplicação para autenticação OAuth.

Criar a base de dados

Crie um novo projeto Lakebase Autoscaling para alojar a sua base de dados. Clique no Comutador de aplicações, selecione Lakebase Postgres, depois crie um novo projeto atribuindo-lhe um nome (por exemplo, my-app-db) e uma versão do Postgres (aceite o Postgres 17 predefinido). Para detalhes completos da configuração, veja Criar um projeto.

Espere que o cálculo fique ativo (aproximadamente 1 minuto) antes de prosseguir.

Passo 2: Configurar autenticação e esquema da base de dados

Crie uma função Postgres para o principal de serviço da sua aplicação com autenticação OAuth, depois crie uma tabela de exemplo com dados para a sua aplicação mostrar.

Configurar autenticação OAuth

No teu projeto Lakebase, abre o Editor SQL e executa estes comandos. A databricks_auth extensão permite a autenticação OAuth. Com esta funcionalidade, as suas funções Postgres aceitam tokens do Databricks em vez de palavras-passe tradicionais.

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

Substitua <DATABRICKS_CLIENT_ID> pelo valor DATABRICKS_CLIENT_ID da sua aplicação. O principal de serviço pode agora autenticar usando tokens OAuth que o Azure Databricks gere automaticamente. Para mais detalhes, veja Criar um papel OAuth para uma identidade Azure Databricks.

Criar esquema de base de dados

Crie uma tabela de exemplo com permissões explícitas para o seu principal de serviço (os principais de serviço não herdam permissões de esquema padrão):

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

Substitui <DATABRICKS_CLIENT_ID> pelo teu DATABRICKS_CLIENT_ID valor.

Passo 3: Construa e configure a sua aplicação

Descarregue os ficheiros da sua aplicação, configure a ligação à base de dados com rotação automática de tokens OAuth e teste localmente antes de implementar.

Descarregar e configurar ficheiros da aplicação

Descarregue os ficheiros da sua aplicação a partir do espaço de trabalho copiando o comando de exportação da secção Sincronizar os ficheiros da aplicação:

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

Edição app.yaml para adicionar detalhes da ligação à base de dados. Obtenha os seus valores de ligação do modal Lakebase Connect selecionando apenas 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>'

Substitua os marcadores:

• <your-endpoint-hostname>: Copiar o valor PGHOST do modal Connect (por exemplo, ep-xyz.database.us-west-2.dev.databricks.com)
• <DATABRICKS_CLIENT_ID>: Utilize o seu DATABRICKS_CLIENT_ID da Etapa 1
• projects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id>: Na aplicação Lakebase, vá ao separador Computes do seu ramo, clique em Obter ID para o seu cálculo e selecione Copiar nome do recurso.

Implementar rotação de tokens OAuth e consulta à base de dados

Substitua app.py por este código, que adiciona rotação automática do token OAuth e uma consulta à base de dados que recolhe as notas do Passo 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)

Os três componentes principais são:

• WorkspaceClient: gera novas credenciais usando o SDK.
• OAuthConnection: uma classe de ligação personalizada que injeta uma nova credencial em cada ligação.
• ConnectionPool: gere as ligações e chama a classe personalizada conforme necessário.

Para mais informações sobre estratégias de rotação de credenciais e gestão de erros, consulte exemplos de rotação de tokens.

Atualização requirements.txt para incluir os pacotes obrigatórios:

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

A versão 0.81.0 ou posterior inclui o generate_database_credential() método.

Teste localmente

Testa a tua aplicação localmente para confirmar se a ligação à base de dados funciona antes de implementar. Ao testar localmente, a aplicação é executada na sua conta de utilizador Azure Databricks (não como a conta de serviço principal), por isso substitua PGUSER pelo seu endereço de email nas variáveis de ambiente abaixo.

Autentique no seu espaço de trabalho e exporte variáveis de ambiente:

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 os valores do seu app.yaml, mas substitua o valor PGUSER (ID do cliente principal de serviço) pelo seu endereço de email Azure Databricks.

Instale dependências e execute a aplicação:

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

Abra http://localhost:8000 no seu navegador. Deverá ver "Olá de Lakebase!" com as suas três notas de exemplo. O pool de ligações gera automaticamente novos tokens OAuth ao criar novas ligações. Para mais detalhes, consulte autenticação de tokens OAuth.

Passo 4: Implementar e verificar

Depois de testares localmente, sincroniza as tuas alterações para uma pasta de workspace e implementa a partir dessa localização:

# 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

Substitua <your-email> pelo seu endereço de email Azure Databricks e <app-name> pelo nome da sua aplicação. O --source-code-path flag indica à implementação para usar os ficheiros carregados em vez da localização padrão da aplicação.

Espere que a implementação termine (2-3 minutos) e depois aceda à sua aplicação no URL fornecido. Deverás ver "Olá da Lakebase!" com as tuas notas de exemplo.

Consulte também

• Exemplos de rotação de tokens
• Aplicações Databricks