Bagikan melalui

Menyambungkan aplikasi eksternal ke Lakebase menggunakan SDK

Penting

Lakebase Autoscaling adalah versi terbaru Lakebase, dengan komputasi penskalaan otomatis, skala-ke-nol, percabangan, dan pemulihan instan. Untuk wilayah yang didukung, lihat Ketersediaan wilayah. Jika Anda adalah pengguna Lakebase Provisioned, lihat Lakebase Provisioned.

Panduan ini menunjukkan cara menghubungkan aplikasi eksternal ke Lakebase Autoscaling menggunakan driver Postgres standar (psycopg, pgx, JDBC) dengan rotasi token OAuth. Anda menggunakan Azure Databricks SDK dengan perwakilan layanan dan kumpulan koneksi yang memanggil generate_database_credential() saat membuka setiap koneksi baru, sehingga Anda mendapatkan token baru (masa pakai 60 menit) setiap kali Anda tersambung. Contoh disediakan untuk Python, Java, dan Go. Untuk penyiapan yang lebih mudah dengan manajemen kredensial otomatis, pertimbangkan Aplikasi Azure Databricks sebagai gantinya.

Apa yang akan Anda buat: Pola koneksi yang menggunakan rotasi token OAuth untuk menyambungkan ke Lakebase Autoscaling dari aplikasi eksternal, lalu memverifikasi koneksi berfungsi.

Anda memerlukan Databricks SDK (Python v0.89.0+, Java v0.73.0+, atau Go v0.109.0+). Selesaikan langkah-langkah berikut secara berurutan:

:::tip Bahasa Lain Untuk bahasa tanpa dukungan Databricks SDK (Node.js, Ruby, PHP, Elixir, Rust, dll.), lihat Menyambungkan aplikasi eksternal ke Lakebase menggunakan API. :::

Cara kerjanya

Databricks SDK menyederhanakan autentikasi OAuth dengan menangani manajemen token ruang kerja secara otomatis:

Alur OAuth SDK

Aplikasi Anda memanggil generate_database_credential() dengan parameter titik akhir. SDK mendapatkan token OAuth ruang kerja secara internal (tidak ada kode yang diperlukan), meminta kredensial database dari Lakebase API, dan mengembalikannya ke aplikasi Anda. Anda kemudian menggunakan kredensial ini sebagai kata sandi saat menyambungkan ke Postgres.

Token OAuth ruang kerja dan kredensial database kedaluwarsa setelah 60 menit. Kumpulan koneksi menangani refresh otomatis dengan memanggil generate_database_credential() saat membuat koneksi baru.

1. Buat prinsipal layanan dengan rahasia OAuth

Buat prinsipal layanan Azure Databricks dengan rahasia OAuth. Detail lengkap ada di Otorisasi akses perwakilan layanan. Untuk membangun aplikasi eksternal, perlu diingat:

  • Atur rahasia Anda ke masa pakai pilihan Anda, hingga 730 hari. Ini menentukan seberapa sering Anda perlu memperbarui kunci rahasia, yang digunakan untuk menghasilkan kredensial database melalui proses rotasi.
  • Aktifkan "Akses ruang kerja" untuk perwakilan layanan (Pengaturan → Identitas dan akses → Perwakilan layanan → {name} tab Konfigurasi →). Hal ini diperlukan untuk menghasilkan kredensial database baru.
  • Perhatikan ID klien (UUID). Anda menggunakannya saat membuat peran Postgres yang cocok dalam penyiapan aplikasi Anda dan untuk PGUSER.

2. Buat peran Postgres untuk prinsipal layanan

UI Lakebase hanya mendukung peran berbasis kata sandi. Buat peran OAuth di Editor Lakebase SQL menggunakan ID klien dari langkah 1 (bukan nama tampilan; nama peran peka huruf besar/kecil):

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

Ganti {client-id} dengan ID klien perwakilan layanan Anda. Lihat Membuat peran OAuth.

3. Dapatkan detail koneksi

Dari proyek Anda di Lakebase Console, klik Sambungkan, pilih cabang dan titik akhir, lalu catat host, database (biasanya databricks_postgres), dan nama endpoint (format: projects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id>).

Atau gunakan CLI:

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

Lihat String koneksi untuk detailnya.

4. Atur variabel lingkungan

Atur variabel lingkungan ini sebelum menjalankan aplikasi Anda:

# 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. Tambahkan kode koneksi

Phyton

Contoh ini menggunakan psycopg3 dengan kelas koneksi kustom yang menghasilkan token baru saat kumpulan membuat setiap koneksi baru.

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

Dependensi:databricks-sdk>=0.89.0, psycopg[binary,pool]>=3.1.0

Go

Contoh ini menggunakan pgxpool dengan panggilan balik BeforeConnect yang menghasilkan token baru untuk setiap koneksi baru.

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

Dependensi: Databricks SDK untuk Go v0.109.0+ (github.com/databricks/databricks-sdk-go), pgx driver (github.com/jackc/pgx/v5)

Catatan: Callback BeforeConnect memastikan token OAuth baru untuk setiap koneksi baru, menangani rotasi token otomatis untuk aplikasi yang berjalan dalam jangka waktu lama.

Java

Contoh ini menggunakan JDBC dengan HikariCP dan DataSource kustom yang menghasilkan token baru saat kumpulan membuat setiap koneksi baru.

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

Dependensi: Databricks SDK untuk Java v0.73.0+ (com.databricks:databricks-sdk-java), driver PostgreSQL JDBC (org.postgresql:postgresql), HikariCP (com.zaxxer:HikariCP)

6. Jalankan dan verifikasi koneksi

Phyton

Pasang dependensi:

pip install databricks-sdk psycopg[binary,pool]

Menjalankan:

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

Output yang diharapkan:

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

Jika current_user cocok dengan ID klien perwakilan layanan Anda dari langkah 1, rotasi token OAuth berfungsi.

Java

Catatan: Ini mengasumsikan Anda memiliki proyek Maven dengan dependensi dari contoh Java di atas di Anda pom.xml.

Pasang dependensi:

mvn install

Menjalankan:

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

Output yang diharapkan:

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

Jika pengguna cocok dengan ID klien prinsipal layanan Anda dari langkah 1, rotasi token OAuth berjalan.

Go

Pasang dependensi:

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

Menjalankan:

go run main.go

Output yang diharapkan:

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

Jika pengguna cocok dengan ID klien prinsipal layanan Anda dari langkah 1, rotasi token OAuth berjalan.

Catatan: Koneksi pertama setelah tidak aktif mungkin memakan waktu lebih lama karena Sistem Penskalaan Otomatis Lakebase mulai komputasi dari awal.

Troubleshooting

Kesalahan Perbaiki
"API dinonaktifkan untuk pengguna tanpa hak akses ruang kerja" Aktifkan "Akses ruang kerja" untuk perwakilan layanan (langkah 1).
"Peran tidak ada" atau autentikasi gagal Buat peran OAuth melalui SQL (langkah 2), bukan UI.
"Koneksi ditolak" atau "Titik akhir tidak ditemukan" Gunakan ENDPOINT_NAME format projects/<id>/branches/<id>/endpoints/<id>; ID titik akhir ada di host.
"Pengguna tidak valid" atau "Pengguna tidak ditemukan" Atur PGUSER ke ID klien perwakilan layanan (UUID), bukan nama tampilan.
  • Last updated on