Ligue uma aplicação externa ao Lakebase usando o SDK

Este guia mostra como ligar aplicações externas ao Autoscaling Lakebase usando drivers Postgres padrão (psycopg, pgx, JDBC) com rotação de token OAuth. Usas o Azure Databricks SDK com um principal de serviço e um pool de ligações que chama generate_database_credential() ao abrir cada nova ligação, por isso recebes um novo token (60 minutos de vida útil) cada vez que te ligas. São fornecidos exemplos para Python, Java e Go. Para uma configuração mais fácil com gestão automática de credenciais, considere o Azure Databricks Apps em vez disso.

O que vais construir: Um padrão de ligação que usa rotação de token OAuth para se ligar ao Autoscaling do Lakebase a partir de uma aplicação externa, e depois verificar se a ligação funciona.

Precisas do SDK Databricks (Python v0.89.0+, Java v0.73.0+, ou Go v0.109.0+). Conclua as seguintes etapas na ordem:

:::dica Outras Linguagens Para linguagens sem suporte ao SDK Databricks (Node.js, Ruby, PHP, Elixir, Rust, etc.), veja como conectar uma aplicação externa ao Lakebase usando a API. :::

Como funciona

O SDK Databricks simplifica a autenticação OAuth ao gerir automaticamente a gestão de tokens do espaço de trabalho:

A sua aplicação chama generate_database_credential() com o parâmetro endpoint. O SDK obtém internamente o token OAuth do workspace (sem necessidade de código), solicita a credencial da base de dados à API do Lakebase e devolve-a à sua aplicação. Depois usa esta credencial como palavra-passe ao ligar-se ao Postgres.

Tanto o token OAuth do workspace como a credencial da base de dados expiram após 60 minutos. As pools de conexões realizam a atualização automática chamando generate_database_credential() ao criar novas ligações.

1. Criar principal de serviço com OAuth secret

Crie uma identidade de serviço do Azure Databricks com um segredo OAuth. Os detalhes completos encontram-se em Autorizar o acesso ao principal do serviço. Para construir uma aplicação externa, tenha em mente:

• Defina o seu segredo para a sua vida preferida, até 730 dias. Isto define com que frequência é necessário atualizar o segredo, que é usado para gerar credenciais de base de dados por rotação.
• Ative "Acesso ao espaço de trabalho" para o principal de serviço (Definições → Identidade e acesso → Principais de serviço → {name} → separador Configurações). É necessário para gerar novas credenciais de base de dados.
• Note o identificador do cliente (um UUID). Utiliza-o ao criar a função Postgres correspondente na configuração da aplicação e para PGUSER.

2. Criar um papel de Postgres para o principal do serviço

Crie um papel OAuth para o principal do serviço. Pode fazer isto na interface do Lakebase (usando o separador do OAuth do diálogo Adicionar função) ou no Editor SQL do Lakebase usando o ID do cliente do passo 1 (não o nome de visualização; o nome da função diferencia maiúsculas de minúsculas):

-- Enable the auth extension (if not already enabled)
CREATE EXTENSION IF NOT EXISTS databricks_auth;

-- Create OAuth role using the service principal client ID
SELECT databricks_create_role('{client-id}', 'SERVICE_PRINCIPAL');

-- Grant database permissions
GRANT CONNECT ON DATABASE databricks_postgres TO "{client-id}";
GRANT USAGE ON SCHEMA public TO "{client-id}";
GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public TO "{client-id}";
ALTER DEFAULT PRIVILEGES IN SCHEMA public
  GRANT SELECT, INSERT, UPDATE, DELETE ON TABLES TO "{client-id}";

Substitua {client-id} pelo seu ID de cliente principal de serviço. Ver Criar funções OAuth.

3. Obter detalhes de ligação

A partir do seu projeto na Lakebase Console, clique em Conectar, selecione branch e endpoint, e anote host, base de dados (normalmente databricks_postgres) e nome do endpoint (formato: projects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id>).

Ou usar a linha de comando:

databricks postgres list-endpoints projects/<project-id>/branches/<branch-id>

Veja cadeias de conexão para mais detalhes.

4. Definir variáveis ambientais

Defina estas variáveis de ambiente antes de executar a sua aplicação:

# Databricks workspace authentication
export DATABRICKS_HOST="https://your-workspace.databricks.com"
export DATABRICKS_CLIENT_ID="<service-principal-client-id>"
export DATABRICKS_CLIENT_SECRET="<your-oauth-secret>"

# Lakebase connection details (from step 3)
export ENDPOINT_NAME="projects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id>"
export PGHOST="<endpoint-id>.database.<region>.cloud.databricks.com"
export PGDATABASE="databricks_postgres"
export PGUSER="<service-principal-client-id>"   # Same UUID as step 1
export PGPORT="5432"
export PGSSLMODE="require"   # Python only

5. Adicionar código de ligação

Python

Este exemplo usa psycopg3 com uma classe de ligação personalizada que gera um token novo quando o pool cria cada nova ligação.

import os
from databricks.sdk import WorkspaceClient
import psycopg
from psycopg_pool import ConnectionPool

# Initialize Databricks SDK
workspace_client = None

def _get_workspace_client():
    """Get or create the workspace client for OAuth."""
    global workspace_client
    if workspace_client is None:
        workspace_client = WorkspaceClient(
            host=os.environ["DATABRICKS_HOST"],
            client_id=os.environ["DATABRICKS_CLIENT_ID"],
            client_secret=os.environ["DATABRICKS_CLIENT_SECRET"],
        )
    return workspace_client

def _get_endpoint_name():
    """Get endpoint name from environment."""
    name = os.environ.get("ENDPOINT_NAME")
    if not name:
        raise ValueError(
            "ENDPOINT_NAME must be set (format: projects/<id>/branches/<id>/endpoints/<id>)"
        )
    return name

class OAuthConnection(psycopg.Connection):
    """Custom connection class that generates a fresh OAuth token per connection."""

    @classmethod
    def connect(cls, conninfo="", **kwargs):
        endpoint_name = _get_endpoint_name()
        client = _get_workspace_client()
        # Generate database credential (tokens are workspace-scoped)
        credential = client.postgres.generate_database_credential(
            endpoint=endpoint_name
        )
        kwargs["password"] = credential.token
        return super().connect(conninfo, **kwargs)

# Create connection pool with OAuth token rotation
def get_connection_pool():
    """Get or create the connection pool."""
    database = os.environ["PGDATABASE"]
    user = os.environ["PGUSER"]
    host = os.environ["PGHOST"]
    port = os.environ.get("PGPORT", "5432")
    sslmode = os.environ.get("PGSSLMODE", "require")

    conninfo = f"dbname={database} user={user} host={host} port={port} sslmode={sslmode}"

    return ConnectionPool(
        conninfo=conninfo,
        connection_class=OAuthConnection,
        min_size=1,
        max_size=10,
        open=True,
    )

# Use the pool in your application
pool = get_connection_pool()
with pool.connection() as conn:
    with conn.cursor() as cur:
        cur.execute("SELECT current_user, current_database()")
        print(cur.fetchone())

Dependências:databricks-sdk>=0.89.0, psycopg[binary,pool]>=3.1.0

Go

Este exemplo usa pgxpool com um callback BeforeConnect que gera um token novo para cada nova ligação.

package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/databricks/databricks-sdk-go"
	"github.com/databricks/databricks-sdk-go/service/postgres"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

func createConnectionPool(ctx context.Context) (*pgxpool.Pool, error) {
	// Initialize Databricks workspace client
	w, err := databricks.NewWorkspaceClient(&databricks.Config{
		Host:         os.Getenv("DATABRICKS_HOST"),
		ClientID:     os.Getenv("DATABRICKS_CLIENT_ID"),
		ClientSecret: os.Getenv("DATABRICKS_CLIENT_SECRET"),
	})
	if err != nil {
		return nil, err
	}

	// Build connection string
	connStr := fmt.Sprintf("host=%s port=%s dbname=%s user=%s sslmode=require",
		os.Getenv("PGHOST"),
		os.Getenv("PGPORT"),
		os.Getenv("PGDATABASE"),
		os.Getenv("PGUSER"))

	config, err := pgxpool.ParseConfig(connStr)
	if err != nil {
		return nil, err
	}

	// Configure pool
	config.MaxConns = 10
	config.MinConns = 1
	config.MaxConnLifetime = 45 * time.Minute
	config.MaxConnIdleTime = 15 * time.Minute

	// Generate fresh token for each new connection
	config.BeforeConnect = func(ctx context.Context, connConfig *pgx.ConnConfig) error {
		credential, err := w.Postgres.GenerateDatabaseCredential(ctx,
			postgres.GenerateDatabaseCredentialRequest{
				Endpoint: os.Getenv("ENDPOINT_NAME"),
			})
		if err != nil {
			return err
		}
		connConfig.Password = credential.Token
		return nil
	}

	return pgxpool.NewWithConfig(ctx, config)
}

func main() {
	ctx := context.Background()
	pool, err := createConnectionPool(ctx)
	if err != nil {
		log.Fatal(err)
	}
	defer pool.Close()

	var user, database string
	err = pool.QueryRow(ctx, "SELECT current_user, current_database()").Scan(&user, &database)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Printf("Connected as: %s to database: %s\n", user, database)
}

Dependências: Databricks SDK para Go v0.109.0+ (github.com/databricks/databricks-sdk-go), Driver pgx (github.com/jackc/pgx/v5)

Nota: O BeforeConnect callback assegura tokens OAuth frescos para cada nova ligação, gerindo a rotação automática dos tokens para aplicações de longa duração.

Java

Este exemplo usa JDBC com HikariCP e uma DataSource personalizada que gera um token novo quando o pool cria cada nova ligação.

import java.sql.*;
import javax.sql.DataSource;
import com.databricks.sdk.WorkspaceClient;
import com.databricks.sdk.core.DatabricksConfig;
import com.databricks.sdk.service.postgres.*;
import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;

public class LakebaseConnection {

    private static WorkspaceClient workspaceClient() {
        String host = System.getenv("DATABRICKS_HOST");
        String clientId = System.getenv("DATABRICKS_CLIENT_ID");
        String clientSecret = System.getenv("DATABRICKS_CLIENT_SECRET");

        return new WorkspaceClient(new DatabricksConfig()
            .setHost(host)
            .setClientId(clientId)
            .setClientSecret(clientSecret));
    }

    private static DataSource createDataSource() {
        WorkspaceClient w = workspaceClient();
        String endpointName = System.getenv("ENDPOINT_NAME");
        String host = System.getenv("PGHOST");
        String database = System.getenv("PGDATABASE");
        String user = System.getenv("PGUSER");
        String port = System.getenv().getOrDefault("PGPORT", "5432");

        String jdbcUrl = "jdbc:postgresql://" + host + ":" + port +
                         "/" + database + "?sslmode=require";

        // DataSource that returns a new connection with a fresh token (tokens are workspace-scoped)
        DataSource tokenDataSource = new DataSource() {
            @Override
            public Connection getConnection() throws SQLException {
                DatabaseCredential cred = w.postgres().generateDatabaseCredential(
                    new GenerateDatabaseCredentialRequest().setEndpoint(endpointName)
                );
                return DriverManager.getConnection(jdbcUrl, user, cred.getToken());
            }

            @Override
            public Connection getConnection(String u, String p) {
                throw new UnsupportedOperationException();
            }
            // ... other DataSource methods (getLogWriter, etc.)
        };

        // Wrap in HikariCP for connection pooling
        HikariConfig config = new HikariConfig();
        config.setDataSource(tokenDataSource);
        config.setMaximumPoolSize(10);
        config.setMinimumIdle(1);
        // Recycle connections before 60-min token expiry
        config.setMaxLifetime(45 * 60 * 1000L);

        return new HikariDataSource(config);
    }

    public static void main(String[] args) throws SQLException {
        DataSource pool = createDataSource();

        try (Connection conn = pool.getConnection();
             Statement st = conn.createStatement();
             ResultSet rs = st.executeQuery("SELECT current_user, current_database()")) {
            if (rs.next()) {
                System.out.println("User: " + rs.getString(1));
                System.out.println("Database: " + rs.getString(2));
            }
        }
    }
}

Dependências: Databricks SDK para Java v0.73.0+ (com.databricks:databricks-sdk-java), driver PostgreSQL JDBC (org.postgresql:postgresql), HikariCP (com.zaxxer:HikariCP)

6. Executar e verificar a ligação

Python

Dependências de instalação:

pip install databricks-sdk psycopg[binary,pool]

Executar:

# Save all the code from step 5 (above) as db.py, then run:
from db import get_connection_pool

pool = get_connection_pool()
with pool.connection() as conn:
    with conn.cursor() as cur:
        cur.execute("SELECT current_user, current_database()")
        print(cur.fetchone())

Resultados esperados:

('c00f575e-d706-4f6b-b62c-e7a14850571b', 'databricks_postgres')

Se current_user corresponder ao ID do cliente principal de serviço do passo 1, a rotação do token OAuth está a funcionar.

Java

Nota: Isto pressupõe que tens um projeto Maven com as dependências do exemplo Java acima no teu pom.xmlarquivo .

Dependências de instalação:

mvn install

Executar:

mvn exec:java -Dexec.mainClass="com.example.LakebaseConnection"

Resultados esperados:

User: c00f575e-d706-4f6b-b62c-e7a14850571b
Database: databricks_postgres

Se o utilizador corresponder ao ID do cliente principal do teu serviço do passo 1, a rotação do token OAuth está a funcionar.

Go

Dependências de instalação:

go mod init myapp
go get github.com/databricks/databricks-sdk-go
go get github.com/jackc/pgx/v5

Executar:

go run main.go

Resultados esperados:

Connected as: c00f575e-d706-4f6b-b62c-e7a14850571b to database: databricks_postgres

Se o utilizador corresponder ao ID do cliente principal do teu serviço do passo 1, a rotação do token OAuth está a funcionar.

Nota: A primeira ligação após um período de inatividade pode demorar mais tempo do que o normal, pois o Autoscaling do Lakebase começa a iniciar do zero.

Solução de problemas