Conectar um aplicativo personalizado do Databricks ao Lakebase

Este tutorial mostra como conectar um aplicativo do Databricks ao Lakebase Autoscaling com rotação automática de credenciais. O aplicativo gera novas credenciais de banco de dados do Databricks antes de expirarem. O exemplo usa o Flask, mas o padrão de autenticação se aplica a qualquer estrutura.

Como funciona

Os Aplicativos do Databricks são autenticados no Lakebase usando tokens OAuth que expiram após uma hora. Para lidar com isso, crie uma função Postgres para a entidade de serviço do aplicativo e configure seu aplicativo para gerar automaticamente tokens novos sempre que precisar se conectar ao banco de dados. O aplicativo usa um pool de conexões. O pool cria novas conexões com tokens novos, conforme necessário, para que seu aplicativo nunca use credenciais expiradas.

Quando você implanta o aplicativo no Azure Databricks, ele é executado como uma entidade de serviço e gera tokens para essa identidade. Quando você testa localmente, o aplicativo é executado como sua conta de usuário Azure Databricks e gera tokens para você. Ambos usam o mesmo código de rotação de token. Somente o contexto de autenticação é alterado.

Antes de começar

Para concluir este tutorial, você precisará:

• Acesso a um workspace do Azure Databricks com dimensionamento automático habilitado para o Lakebase Postgres. Entre em contato com o administrador do workspace se você não vir o Lakebase no comutador de aplicativos.
• Permissão para criar aplicativos
• Familiaridade básica com Python e SQL
• CLI do Databricks instalada para desenvolvimento local
• Python 3.9 ou posterior instalado localmente

Etapa 1: Criar seu aplicativo e banco de dados

Primeiro, crie um aplicativo databricks e um projeto lakebase. O aplicativo obtém automaticamente uma identidade de entidade de serviço que você usará para autenticação de banco de dados.

Criar o aplicativo

Crie um novo aplicativo Databricks usando o modelo Hello world do Flask. Consulte Criar um aplicativo do Databricks a partir de um modelo.

Depois de criar o aplicativo, vá para a guia Ambiente do aplicativo e anote o DATABRICKS_CLIENT_ID valor (formato UUID como 6b215d2b-f099-4bdb-900a-60837201ecec). Isso se torna o nome de usuário do Postgres do aplicativo para autenticação OAuth.

Criar o banco de dados

Crie um novo projeto de dimensionamento automático do Lakebase para hospedar seu banco de dados. Clique no no App Switcher, selecione Lakebase Postgres e, em seguida, crie um novo projeto com um nome (por exemplo my-app-db) e uma versão do Postgres (aceite o Postgres 17 padrão). Para obter detalhes completos da instalação, consulte Criar um projeto.

Aguarde até que a computação se torne ativa (aproximadamente 1 minuto) antes de continuar.

Etapa 2: Configurar a autenticação e o esquema do banco de dados

Crie um papel Postgres para o principal de serviço do aplicativo com autenticação OAuth e, em seguida, crie uma tabela de exemplo com dados para o aplicativo exibir.

Configurar a autenticação OAuth

Em seu projeto do Lakebase, abra o Editor de SQL e execute esses comandos. A databricks_auth extensão habilita a autenticação OAuth. Com ele, suas funções do Postgres aceitam tokens do Databricks em vez de senhas 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 do DATABRICKS_CLIENT_ID aplicativo. A principal de serviço agora pode se autenticar usando tokens OAuth que o Azure Databricks gerencia automaticamente. Para obter detalhes, consulte Criar uma função OAuth para uma identidade Azure Databricks.

Criar esquema de banco de dados

Crie uma tabela de exemplo com permissões explícitas para sua entidade de serviço (as entidades 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');

Substitua <DATABRICKS_CLIENT_ID> pelo seu DATABRICKS_CLIENT_ID valor.

Etapa 3: Criar e configurar seu aplicativo

Baixe os arquivos do aplicativo, configure a conexão de banco de dados com a rotação automática do token OAuth e teste localmente antes da implantação.

Baixar e configurar arquivos de aplicativo

Baixe os arquivos do aplicativo do workspace copiando o comando de exportação da seção Sincronizar os arquivos do aplicativo:

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

Edite app.yaml para adicionar detalhes de conexão de banco de dados. Obtenha os valores de conexão do modal do Lakebase Connect selecionando somente 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 espaços reservados:

• <your-endpoint-hostname>: copie o valor PGHOST do modal Connect (por exemplo, ep-xyz.database.us-west-2.dev.databricks.com)
• <DATABRICKS_CLIENT_ID>: Use o seu DATABRICKS_CLIENT_ID da Etapa 1
• projects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id>: no aplicativo Lakebase, acesse a guia Computação da sua ramificação, clique em Obter ID para sua computação e selecione Copiar nome do recurso.

Implementar a rotação do token OAuth e a consulta de banco de dados

Substitua app.py por esse código, que adiciona rotação automática de token OAuth e uma consulta de banco de dados que busca as anotações da Etapa 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 principais componentes são:

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

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

Atualize requirements.txt para incluir os pacotes necessá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.

Testar localmente

Teste seu aplicativo localmente para verificar se a conexão de banco de dados funciona antes da implantação. Ao testar localmente, o aplicativo é executado com a sua conta de usuário do Azure Databricks (e não o principal do serviço), então altere PGUSER para o seu endereço de e-mail nas variáveis de ambiente abaixo.

Autentique-se no workspace 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 de seu app.yaml, mas substitua o valor PGUSER (ID do cliente do service principal) pelo seu endereço de email da Azure Databricks.

Instale dependências e execute o aplicativo:

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

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

Etapa 4: Implantar e verificar

Depois de testar localmente, sincronize suas alterações em uma pasta de workspace e implante a partir desse local:

# 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 do aplicativo. O --source-code-path flag instrui a implantação a usar seus arquivos carregados em vez do local padrão do aplicativo.

Aguarde a conclusão da implantação (de 2 a 3 minutos) e acesse seu aplicativo na URL fornecida. Você deve ver "Olá de Lakebase!" junto com suas notas de exemplo.

Consulte também

• Exemplos de rotação de token
• Aplicativos do Databricks