Připojení externí aplikace k Lakebase pomocí sady SDK

Tato příručka ukazuje, jak připojit externí aplikace k automatickému škálování Lakebase pomocí standardních ovladačů Postgres (psycopg, pgx, JDBC) s obměnou tokenů OAuth. Používáte Azure Databricks SDK se služebním principálem a fondem připojení, který volá generate_database_credential() při otevírání každého nového připojení, takže při každém připojení získáte nový token (s 60 minutovou životností). Příklady jsou k dispozici pro Python, Javu a Go. Pro snadnější nastavení pomocí automatické správy přihlašovacích údajů zvažte místo toho Azure Databricks Apps .

Co vytvoříte: Model připojení, který používá obměnu tokenů OAuth pro připojení k automatickému škálování Lakebase z externí aplikace, a pak ověřte, že připojení funguje.

Potřebujete sadu Databricks SDK (Python v0.89.0+, Java v0.73.0 nebo Go v0.109.0+). Proveďte následující kroky v pořadí:

:::tip Další jazyky pro jazyky bez podpory sady Databricks SDK (Node.js, Ruby, PHP, Elixir, Rust atd.), viz Připojení externí aplikace k Lakebase pomocí rozhraní API. :::

Jak to funguje

Sada Databricks SDK zjednodušuje ověřování OAuth automatickým zpracováním správy tokenů pracovního prostoru:

OAuth tok SDK

Vaše aplikace volá generate_database_credential() pomocí parametru koncového bodu. Sada SDK získá token OAuth pracovního prostoru interně (nevyžaduje žádný kód), vyžádá si přihlašovací údaje databáze z rozhraní API Lakebase a vrátí ho do vaší aplikace. Tyto přihlašovací údaje pak použijete jako heslo při připojování k Postgresu.

Platnost tokenu OAuth pracovního prostoru i přihlašovacích údajů databáze vyprší po 60 minutách. Fondy připojení zpracovávají automatickou aktualizaci voláním generate_database_credential() při vytváření nových připojení.

1. Vytvoření instančního objektu s tajným kódem OAuth

Vytvořte instanční objekt Azure Databricks s tajným kódem OAuth. Úplné podrobnosti jsou v Autorizace přístupu služebního objektu. Při vytváření externí aplikace mějte na paměti:

  • Nastavte tajný kód na upřednostňovanou dobu života až do 730 dnů. To definuje, jak často potřebujete aktualizovat tajný klíč, který se používá k vygenerování přihlašovacích údajů databáze prostřednictvím obměny.
  • Pro objekt služby povolte přístup k pracovnímu prostoru (v nastavení → Identita a přístup → Služební principálové → na kartě Konfigurace). Vyžaduje se pro generování nových přihlašovacích údajů databáze.
  • Poznamenejte si ID klienta (UUID). Použijete ji při vytváření odpovídající role Postgres v nastavení aplikace a pro PGUSER.

2. Vytvoření role Postgres pro služební principal

Vytvořte roli OAuth pro služební principál. Můžete to udělat v uživatelském rozhraní Lakebase (pomocí karty OAuth dialogového okna Přidat roli) nebo v Editoru Lakebase SQL pomocí ID klienta z kroku 1 (ne zobrazovaný název, název role se rozlišují malá a velká písmena):

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

Nahraďte {client-id} ID klienta služby principal. Viz Vytvoření rolí OAuth.

3. Získání podrobností o připojení

V projektu v konzole Lakebase klikněte na Připojit, vyberte větev a koncový bod a poznamenejte si hostitele, databázi (obvykle databricks_postgres) a název koncového bodu (formát: projects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id>).

Nebo použijte rozhraní příkazového řádku:

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

Podrobnosti najdete v části Připojovací řetězce .

4. Nastavení proměnných prostředí

Před spuštěním aplikace nastavte tyto proměnné prostředí:

# 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. Přidání kódu připojení

Python

Tento příklad používá psycopg3 s vlastní třídou připojení, která generuje nový token, když fond vytvoří každé nové připojení.

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

Závislosti:databricks-sdk>=0.89.0, psycopg[binary,pool]>=3.1.0

Go

Tento příklad používá pgxpool se zpětným voláním BeforeConnect, který generuje nový token pro každé nové připojení.

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)
}

Závislosti: Databricks SDK pro Go v0.109.0+ (github.com/databricks/databricks-sdk-go), ovladač pgx (github.com/jackc/pgx/v5)

Poznámka:BeforeConnect Zpětné volání zajišťuje nové tokeny OAuth pro každé nové připojení a zpracovává automatické obměny tokenů pro dlouhotrvající aplikace.

Java

Tento příklad používá JDBC s HikariCP a vlastní Zdroj dat, který generuje nový token, když fond vytvoří každé nové připojení.

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));
            }
        }
    }
}

Závislosti: Databricks SDK pro Javu v0.73.0+ (com.databricks:databricks-sdk-java), ovladač PostgreSQL JDBC (org.postgresql:postgresql), HikariCP (com.zaxxer:HikariCP)

6. Spusťte a ověřte připojení.

Python

Nainstalujte závislosti:

pip install databricks-sdk psycopg[binary,pool]

Spustit:

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

Očekávaný výstup:

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

Pokud current_user odpovídá ID klienta principála služby z kroku 1, je obměna tokenu OAuth funkční.

Java

Poznámka: Předpokládá se, že máte projekt Maven se závislostmi z výše uvedeného příkladu Java ve vašem pom.xml.

Nainstalujte závislosti:

mvn install

Spustit:

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

Očekávaný výstup:

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

Pokud se uživatel shoduje s ID klienta hlavního objektu služby z kroku 1, funguje obměna tokenu OAuth.

Go

Nainstalujte závislosti:

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

Spustit:

go run main.go

Očekávaný výstup:

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

Pokud se uživatel shoduje s ID klienta hlavního objektu služby z kroku 1, funguje obměna tokenu OAuth.

Poznámka: První připojení po nečinnosti může trvat déle, protože automatické škálování LakeBase spouští výpočetní prostředky od nuly.

Řešení problémů

Error Oprava
Rozhraní API je pro uživatele bez oprávnění přístupu k pracovnímu prostoru deaktivováno. Povolte "přístup k pracovnímu prostoru" pro službu (krok 1).
Role neexistuje nebo se ověření nezdaří Vytvořte roli OAuth prostřednictvím SQL (krok 2), ne uživatelského rozhraní.
"Připojení odmítnuto" nebo "Koncový bod nebyl nalezen" Použijte ENDPOINT_NAME formát projects/<id>/branches/<id>/endpoints/<id>. ID koncového bodu je v host.
Neplatný uživatel nebo Uživatel nebyl nalezen Nastavte PGUSER na ID klienta principála služby (UUID), nikoli na jeho zobrazovaný název.
  • Last updated on