Connecter une application externe à Lakebase à l'aide du SDK
Ce guide montre comment connecter des applications externes à Lakebase Autoscaling à l'aide de drivers Postgres standards (psycopg, pgx, JDBC) avec rotation de jeton OAuth. Vous utilisez le SDK Databricks avec un Service Principal et un pool de connexions qui appelle generate_database_credential() à l'ouverture de chaque nouvelle connexion, ce qui vous permet d'obtenir un nouveau jeton (durée de vie de 60 minutes) chaque fois que vous vous connectez. Des exemples sont fournis pour Python, Java et Go. Pour une configuration plus facile avec la gestion automatique des informations d'identification, envisagez plutôt Databricks Apps.
Ce que vous allez créer : un modèle de connexion qui utilise la rotation des jetons OAuth pour se connecter à la mise à l'échelle automatique Lakebase à partir d'une application externe, puis vérifier que la connexion fonctionne.
Vous avez besoin du Databricks SDK (Python v0.89.0+, Java v0.73.0+, ou Go v0.109.0+). Effectuez les étapes suivantes dans l'ordre :
:::tip Autres langages Pour les langages sans support du SDK Databricks (Node.js, Ruby, PHP, Elixir, Rust, etc.), voir Connecter une application externe à Lakebase à l'aide de l'API. :::
Comment cela fonctionne
Le SDK Databricks simplifie l'authentification OAuth en gérant automatiquement la gestion des jetons du Workspace :
Votre application appelle generate_database_credential() avec le paramètre d’endpoint. Le SDK obtient le jeton OAuth du Workspace en interne (aucun code requis), demande l’identifiant de la base de données à l’API Lakebase et le renvoie à votre application. Vous utilisez ensuite cet identifiant comme mot de passe lors de la connexion à Postgres.
Le jeton OAuth du Workspace et les identifiants de base de données expirent tous deux après 60 minutes. Les pools de connexions gèrent l' refresh automatique en appelant generate_database_credential() lors de la création de nouvelles connexions.
1. Créez un Service Principal avec un secret OAuth
Créez un Service Principal Databricks avec un secret OAuth. Tous les détails se trouvent dans Autoriser l'accès au Service Principal. Pour créer une application externe, gardez à l’esprit :
- Définissez votre secret avec la durée de vie que vous préférez, jusqu'à 730 jours. Ceci définit la fréquence à laquelle vous devez refresh le secret, qui est utilisé pour générer des identifiants de base de données via la rotation.
- Activez « l'accès au Workspace » pour le Service Principal (Paramètres → Identité et accès → Service Principals →
{name}→ tab Configurations). Ceci est requis pour générer de nouveaux identifiants de base de données. - Notez l'ID client (un UUID). Vous l'utilisez lors de la création du rôle Postgres correspondant dans la configuration de votre application et pour
PGUSER.
2. Créer un rôle Postgres pour le Service Principal
Créez un rôle OAuth pour le Service Principal. Vous pouvez le faire dans l’interface utilisateur de Lakebase (à l’aide de la **tab** OAuth de la boîte de dialogue **Ajouter un rôle**) ou dans l’éditeur SQL de Lakebase en utilisant l’ID client de l’étape 1 (pas le nom d’affichage ; le nom du rôle est sensible à la casse) :
-- 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}";
Remplacez {client-id} par l'ID client de votre Service Principal. Voir Créer des rôles OAuth.
3. Obtenir les détails de la connexion
Depuis votre projet dans la console Lakebase, cliquez sur **Connecter**, sélectionnez la Branch et l'Endpoint, et notez l'**hôte**, la **base de données** databricks_postgres (généralement) et le **nom de l'Endpoint** (formatprojects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id>:).
Ou utilisez la CLI :
databricks postgres list-endpoints projects/<project-id>/branches/<branch-id>
Consultez les chaînes de connexion pour plus de détails.
4. Définir les variables d'environnement
Définissez ces variables d'environnement avant d'exécuter votre application :
# 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. Ajouter du code de connexion
- Python
- Go
- Java
Cet exemple utilise psycopg3 avec une classe de connexion personnalisée qui génère un nouveau jeton lorsque le pool crée chaque nouvelle connexion.
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())
Dépendances : databricks-sdk>=0.89.0, psycopg[binary,pool]>=3.1.0
Cet exemple utilise pgxpool avec un callback BeforeConnect qui génère un nouveau jeton pour chaque nouvelle connexion.
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)
}
Dépendances : Databricks SDK for Go v0.109.0+ (github.com/databricks/databricks-sdk-go), driver pgx (github.com/jackc/pgx/v5)
Remarque : la fonction de rappel BeforeConnect assure des jetons OAuth actualisés pour chaque nouvelle connexion, gérant la rotation automatique des jetons pour les applications à exécution prolongée.
Cet exemple utilise JDBC avec HikariCP et une source de données personnalisée qui génère un jeton frais lorsque le Pool crée chaque nouvelle connexion.
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));
}
}
}
}
Dépendances : Databricks SDK pour Java v0.73.0+ (),com.databricks:databricks-sdk-java PostgreSQL JDBC Driverorg.postgresql:postgresql (), HikariCP ().com.zaxxer:HikariCP
6. Exécuter et vérifier la connexion
- Python
- Java
- Go
Installer les dépendances :
pip install databricks-sdk psycopg[binary,pool]
Exécuter :
# 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())
Sortie attendue :
('c00f575e-d706-4f6b-b62c-e7a14850571b', 'databricks_postgres')
Si current_user correspond à l'ID client de votre Service Principal de l'étape 1, la rotation du jeton OAuth fonctionne.
Remarque : Ceci suppose que vous avez un projet Maven avec les dépendances de l'exemple Java ci-dessus dans votre pom.xml.
Installer les dépendances :
mvn install
Exécuter :
mvn exec:java -Dexec.mainClass="com.example.LakebaseConnection"
Sortie attendue :
User: c00f575e-d706-4f6b-b62c-e7a14850571b
Database: databricks_postgres
Si l'utilisateur correspond à l'ID client de votre Service Principal de l'étape 1, la rotation du jeton OAuth fonctionne.
Installer les dépendances :
go mod init myapp
go get github.com/databricks/databricks-sdk-go
go get github.com/jackc/pgx/v5
Exécuter :
go run main.go
Sortie attendue :
Connected as: c00f575e-d706-4f6b-b62c-e7a14850571b to database: databricks_postgres
Si l'utilisateur correspond à l'ID client de votre Service Principal de l'étape 1, la rotation du jeton OAuth fonctionne.
Note: La première connexion après inactivité peut prendre plus de temps, car le dimensionnement automatique de Lakebase start le compute à partir de zéro.
Dépannage
Erreur | Corriger |
|---|---|
« L'API est désactivée pour les utilisateurs sans droit d'accès au workspace » | Activez "Accès au Workspace" pour le Service Principal (étape 1). |
« Le rôle n'existe pas » ou l'authentification échoue | Créez le rôle OAuth via SQL (étape 2), et non l'interface utilisateur. |
« Connexion refusée » ou « Endpoint introuvable » | Utilisez le format |
"Utilisateur non valide" ou "Utilisateur introuvable" | Définissez |