Conecte aplicativos externos ao Lakebase usando o SDK.

info

O dimensionamento automático do Lakebase está disponível nas seguintes regiões: us-east-1, us-east-2, us-west-2, eu-central-1, eu-west-1, 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 guia mostra como conectar aplicativos externos ao recurso de escalonamento automático do Lakebase usando drivers padrão do Postgres (psycopg, pgx, JDBC) com rotação de tokens OAuth . Você usa o SDK Databricks com uma entidade de serviço e um pool de conexões que chama generate_database_credential() ao abrir cada nova conexão, então você obtém novos tokens (com duração de 60 minutos) cada vez que se conecta. São fornecidos exemplos para Python, Java e Go. Para uma configuração mais fácil com gerenciamento automático de credenciais, considere usar o Databricks Apps .

O que você vai construir: Um padrão de conexão que usa a rotação de tokens OAuth para conectar-se ao Lakebase Autoscale a partir de um aplicativo externo e, em seguida, verificar se a conexão funciona.

Você precisa do SDK do Databricks (Python v0.89.0 ou superior). Java v0.73.0+ ou Go v0.109.0+). Conclua os seguintes passos em ordem:

:::dica Outros idiomas Para linguagens sem suporte ao SDK do Databricks (Node.js, Ruby, PHP, Elixir, Rust, etc.), veja Conectar aplicativo externo ao Lakebase usando API. :::

nota

Escopo de tokens: tokens de credenciais do banco de dados têm escopo de workspace . Embora o parâmetro endpoint seja obrigatório, os tokens retornados podem acessar qualquer banco de dados ou projeto no workspace para o qual a entidade de serviço tenha permissões.

Como funciona

O SDK Databricks simplifica a autenticação OAuth ao gerenciar automaticamente os tokens workspace :

Fluxo OAuth do SDK

Seu aplicativo chama generate_database_credential() com o parâmetro de endpoint. O SDK obtém os tokens OAuth workspace internamente (sem necessidade de código), solicita as credenciais do banco de dados da API do Lakebase e as retorna para o seu aplicativo. Em seguida, você usa essa credencial como senha ao se conectar ao Postgres.

Tanto os tokens OAuth workspace quanto as credenciais do banco de dados expiram após 60 minutos. O pool de conexões lida com refresh automática chamando generate_database_credential() ao criar novas conexões.

1. Crie uma entidade de serviço com um segredo OAuth

Crie uma entidade de serviço do Databricks com um segredo OAuth. Todos os detalhes estão em Autorizar acesso da entidade de serviço. Para criar um aplicativo externo, tenha em mente:

Defina a duração do seu segredo como desejar, até 730 dias. Isso define a frequência com que você precisa refresh o segredo, que é usado para gerar credenciais de banco de dados por meio de rotação.
Habilite o "acesso ao espaço de trabalho" para a entidade de serviço (Configurações → Identidade e acesso → entidade de serviço → {name} → tab Configurações). É necessário para gerar novas credenciais de banco de dados.
Anote o ID do cliente (um UUID). Você o utiliza ao criar a função Postgres correspondente na configuração do seu aplicativo e para PGUSER.

2. Crie uma função do Postgres para a entidade de serviço.

A interface de usuário do Lakebase suporta apenas funções baseadas em senha. Crie uma função OAuth no Editor SQL do Lakebase usando o ID do cliente da etapa 1 (não o nome de exibição; o nome da função diferencia maiúsculas de minúsculas):

SQL
-- 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 ID do cliente da sua entidade de serviço. Consulte Criar funções OAuth.

3. Obtenha os detalhes da conexão

No seu projeto no Console do Lakebase, clique em Conectar , selecione branch e endpoint e anote host , banco de dados (geralmente databricks_postgres) e nome do endpoint (formato: projects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id>).

Ou utilize a CLI:

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

Consulte stringsde conexão para obter detalhes.

4. Defina a variável de ambiente

Defina estas variáveis de ambiente antes de executar seu aplicativo:

Bash
# 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>.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. Adicione o código de conexão.

Python
Go
Java

Este exemplo utiliza o psycopg3 com uma classe de conexão personalizada que gera novos tokens sempre que o pool cria uma nova conexão.

Python
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

Este exemplo utiliza o pgxpool com um callback BeforeConnect que gera novos tokens para cada nova conexão.

Go
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: SDK do Databricks para Go v0.109.0+ (github.com/databricks/databricks-sdk-go), driver pgx (github.com/jackc/pgx/v5)

Nota: O callback BeforeConnect garante tokens OAuth atualizados para cada nova conexão, gerenciando a rotação automática de tokens para aplicações de longa duração.

Este exemplo utiliza JDBC com HikariCP e uma fonte de dados personalizada que gera novos tokens sempre que o pool cria uma nova conexão.

Java
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 JDBC do PostgreSQL (org.postgresql:postgresql), HikariCP (com.zaxxer:HikariCP)

6. execução e verificação da conexão

Python
Java
Go

Instalar dependências:

Bash
pip install databricks-sdk psycopg[binary,pool]

execução:

Python
# 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())

Resultado esperado:

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

Se current_user corresponder ao ID do cliente da entidade de serviço da etapa 1, a rotação de tokens OAuth está funcionando.

Nota: Isto pressupõe que você tenha um projeto Maven com as dependências do exemplo Java acima em seu pom.xml.

Instalar dependências:

Bash
mvn install

execução:

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

Resultado esperado:

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

Se o usuário corresponder ao ID do cliente da sua entidade de serviço da etapa 1, a rotação de tokens OAuth está funcionando.

Instalar dependências:

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

execução:

Bash
go run main.go

Resultado esperado:

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

Se o usuário corresponder ao ID do cliente da sua entidade de serviço da etapa 1, a rotação de tokens OAuth está funcionando.

Observação: A primeira conexão após o Parado pode demorar mais, pois o dimensionamento automático do Lakebase inicia compute a partir do zero.

Solução de problemas

Erro	Consertar
"API está desativada para usuários sem direito de acesso workspace ."	Habilite o "acesso ao espaço de trabalho" para a entidade de serviço (o passo 1).
"A função não existe" ou a autenticação falhou	Crie a função OAuth via SQL (passo 2), não pela interface do usuário.
"Conexão recusada" ou "Ponto de extremidade não encontrado"	Use o formato `ENDPOINT_NAME` `projects/<id>/branches/<id>/endpoints/<id>`; o ID do endpoint está no host.
"Usuário inválido" ou "Usuário não encontrado"	Defina `PGUSER` como o ID do cliente da entidade de serviço (UUID), e não o nome de exibição.

Como funciona​

1. Crie uma entidade de serviço com um segredo OAuth​

2. Crie uma função do Postgres para a entidade de serviço.​

3. Obtenha os detalhes da conexão​

4. Defina a variável de ambiente​

5. Adicione o código de conexão.​

6. execução e verificação da conexão​

Solução de problemas​

Como funciona

1. Crie uma entidade de serviço com um segredo OAuth

2. Crie uma função do Postgres para a entidade de serviço.

3. Obtenha os detalhes da conexão

4. Defina a variável de ambiente

5. Adicione o código de conexão.

6. execução e verificação da conexão

Solução de problemas