Tutorial: Conecte um aplicativo ao dimensionamento automático do Lakebase

info

O dimensionamento automático do Lakebase está disponível nas seguintes regiões: us-east-1, us-east-2, eu-central-1, eu-west-1, eu-west-2, ap-south-1, ap-southeast-1, ap-southeast-2.

O Lakebase autoscale é a versão mais recente do Lakebase com recursos como autoscale compute, escala-to-zero, branching e instant restore. Para comparação de recursos com o provisionamento do Lakebase, veja escolhendo entre versões.

Este tutorial mostra como conectar um aplicativo Databricks ao Lakebase com escalonamento automático e rotação automática de credenciais. O aplicativo gera novas credenciais de banco de dados do Databricks antes que elas expirem. O exemplo utiliza Flask , mas o padrão de autenticação se aplica a qualquer framework.

Como funciona

Databricks Apps se autenticam no Lakebase usando tokens OAuth que expiram após uma hora. Para lidar com isso, você cria uma função do Postgres para a entidade de serviço do seu aplicativo e, em seguida, configura seu aplicativo para gerar automaticamente novos tokens sempre que precisar se conectar ao banco de dados. Isso acontece por meio de um padrão de pool de conexões: o pool cria novas conexões com tokens atualizados conforme necessário, para que seu aplicativo nunca use credenciais expiradas.

Quando você implanta o aplicativo no Databricks, ele executa como sua entidade de serviço e gera tokens para essa identidade. Ao testar localmente, a execução do aplicativo utiliza sua account de usuário Databricks e gera tokens para você. Ambos usam o mesmo código de rotação de tokens: apenas o contexto de autenticação muda.

Antes de começar

Para concluir este tutorial, você precisa de:

• Acesso a um workspace Databricks com o recurso de escalonamento automático do Lakebase Postgres ativado (entre em contato com o administrador do seu workspace se você não vir o Lakebase no seletor de aplicativos).
• Permissão para criar Databricks Apps
• Conhecimento básico de Python e SQL
• CLI do Databricks instalada para desenvolvimento local.
• Python 3.9 ou posterior instalado localmente

Passo 1: Crie 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 do banco de dados.

Crie o aplicativo

Crie um novo aplicativo Databricks usando o padrão Flask Hello World ( + Novo > Aplicativo em seu workspace). Para obter instruções detalhadas, consulte Criar um aplicativo Databricks.

Após a instalação, acesse a tab Ambiente do aplicativo e anote o valor DATABRICKS_CLIENT_ID (formato UUID como 6b215d2b-f099-4bdb-900a-60837201ecec). Este será o nome de usuário do Postgres do seu aplicativo para autenticação OAuth.

nota

Não implante o aplicativo ainda - você configurará a conexão com o banco de dados primeiro.

Criar o banco de dados

Crie um novo projeto de escalonamento automático do Lakebase para hospedar seu banco de dados. Clique No canto superior direito, selecione Aplicativos , selecione Lakebase e crie um novo projeto com o nome desejado (por exemplo, my-app-db) e a versão do Postgres (aceite o default Postgres 17). Para obter detalhes completos sobre a configuração, consulte Criar um projeto.

Aguarde até que o compute fique ativo (aproximadamente 1 minuto) antes de prosseguir.

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

Crie uma função do Postgres para a entidade de serviço do seu aplicativo com autenticação OAuth e, em seguida, crie uma tabela de exemplo com os dados que seu aplicativo exibirá.

Configurar autenticação OAuth

No seu projeto Lakebase, abra o Editor SQL e execute estes comandos. A extensão databricks_auth habilita a autenticação OAuth, permitindo que suas funções do Postgres aceitem tokens do Databricks em vez de senhas tradicionais:

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

Substitua <DATABRICKS_CLIENT_ID> pelo valor DATABRICKS_CLIENT_ID do seu aplicativo. A entidade de serviço agora pode autenticar usando tokens OAuth que Databricks gerencia automaticamente. Para obter detalhes, consulte Criar uma função OAuth para uma identidade Databricks.

Criar esquema de banco de dados

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

Substitua <DATABRICKS_CLIENT_ID> pelo seu valor DATABRICKS_CLIENT_ID .

a etapa 3: Crie e configure seu aplicativo

Baixe os arquivos do seu aplicativo, configure a conexão com o banco de dados com a rotação automática de tokens OAuth e teste localmente antes de implantá-lo.

Baixe e configure os arquivos do aplicativo.

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

Bash

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 com o banco de dados. Obtenha os valores de conexão na janela modal do Lakebase Connect selecionando "Somente parâmetros" :

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

Substitua os marcadores de posição:

• <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 seu DATABRICKS_CLIENT_ID do passo 1
• projects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id> : execução databricks postgres list-endpoints e copie o valor do campo name (formato: projects/<uuid>/branches/<uuid>/endpoints/<id>)

Implementar rotação de tokens OAuth e consulta de banco de dados

Substitua app.py por este código. Esta implementação introduz duas alterações key : rotação automática de tokens OAuth para autenticação e uma consulta à base de dados que recupera e exibe as notas que você criou no Passo 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
        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)

Este código implementa a rotação automática de credenciais de banco de dados usando três componentes key :

• WorkspaceClient : Gera novas credenciais de banco de dados por meio do SDK do Databricks.
• Classe OAuthConnection : Classe de conexão personalizada que injeta uma nova credencial em cada nova conexão.
• ConnectionPool : gerencia conexões de banco de dados e chama a classe de conexão personalizada ao criar conexões.

Esse padrão garante que seu aplicativo sempre tenha credenciais de banco de dados válidas. Para uma explicação detalhada de como funciona a rotação de credenciais, diferentes estratégias de rotação e tratamento de erros , consulte Rotação de tokens no Lakebase.

Atualize requirements.txt para incluir o pacote necessário:

Txt

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

A versão mínima SDK (0.81.0) garante que o método generate_database_credential() esteja disponível para geração de tokens OAuth .

Teste localmente

Teste seu aplicativo localmente para verificar se a conexão com o banco de dados funciona antes da implantação. Ao testar localmente, o aplicativo é executado como sua account de usuário Databricks (não a entidade de serviço), portanto, você precisará alterar PGUSER para seu endereço email na variável de ambiente abaixo.

Autentique-se em seu workspace e exporte variável de ambiente:

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

Copie os valores de app.yaml, mas substitua o valor de PGUSER (ID do cliente da entidade de serviço) pelo seu endereço email Databricks .

Instale dependências e execute o aplicativo:

Bash

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

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

o passo 4: implantado e verificado

Após realizar os testes locais, sincronize as alterações com uma pasta workspace e implante a partir desse local:

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

Substitua <your-email> pelo seu endereço email Databricks e <app-name> pelo nome do seu aplicativo. A flag --source-code-path indica à implantação que use os seus arquivos de upload em vez do local default do aplicativo.

Aguarde a conclusão da implantação (2 a 3 minutos) e, em seguida, acesse seu aplicativo no URL fornecido. Você deverá ver "Olá do Lakebase!" junto com suas anotações de exemplo.

Veja também

• Autenticação no Lakebase
• Documentação Databricks Apps