Menyambungkan aplikasi Databricks kustom ke Lakebase

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.

Tutorial ini menunjukkan kepada Anda cara menyambungkan aplikasi Databricks ke Lakebase Autoscaling dengan rotasi kredensial otomatis. Aplikasi ini menghasilkan kredensial database baru dari Databricks sebelum kedaluwarsa. Contohnya menggunakan Flask, tetapi pola autentikasi berlaku untuk kerangka kerja apa pun.

Cara kerjanya

Aplikasi Databricks mengautentikasi ke Lakebase menggunakan token OAuth yang kedaluwarsa setelah satu jam. Untuk menangani hal ini, Anda membuat peran Postgres untuk perwakilan layanan aplikasi Anda, lalu mengonfigurasi aplikasi Anda untuk secara otomatis menghasilkan token baru setiap kali perlu terhubung ke database. Aplikasi ini menggunakan kumpulan koneksi. Kumpulan membuat koneksi baru dengan token baru sesuai kebutuhan, sehingga aplikasi Anda tidak pernah menggunakan kredensial yang kedaluwarsa.

Saat Anda menyebarkan aplikasi ke Azure Databricks, aplikasi berjalan sebagai perwakilan layanannya dan menghasilkan token untuk identitas tersebut. Saat Anda menguji secara lokal, aplikasi berjalan sebagai akun pengguna Azure Databricks Anda dan menghasilkan token untuk Anda. Keduanya menggunakan kode rotasi token yang sama. Hanya konteks autentikasi yang berubah.

Sebelum Anda memulai

Untuk menyelesaikan tutorial ini, Anda memerlukan:

  • Akses ke ruang kerja Azure Databricks dengan fitur autoscaling Lakebase Postgres yang diaktifkan. Hubungi admin ruang kerja Jika Anda tidak melihat Lakebase di pengalih aplikasi.
  • Izin untuk membuat aplikasi
  • Pemahaman dasar tentang Python dan SQL
  • Databricks CLI diinstal untuk pengembangan lokal
  • Python 3.9 atau yang lebih baru diinstal secara lokal

Langkah 1: Membuat aplikasi dan database Anda

Pertama, buat aplikasi Databricks dan proyek Lakebase. Aplikasi ini secara otomatis mendapatkan identitas principal layanan yang akan Anda gunakan untuk otentikasi basis data.

Membuat aplikasi

Buat aplikasi Databricks baru menggunakan templat Flask Hello world . Lihat Membuat aplikasi Databricks dari templat.

Setelah Anda membuat aplikasi, buka tab Lingkungan aplikasi dan catat nilainya DATABRICKS_CLIENT_ID (format UUID seperti 6b215d2b-f099-4bdb-900a-60837201ecec). Ini menjadi nama pengguna Postgres aplikasi Anda untuk autentikasi OAuth.

Nota

Jangan meluncurkan aplikasi dulu. Konfigurasikan koneksi database terlebih dahulu.

Membuat database

Buat proyek Lakebase Autoscaling baru untuk menghosting database Anda. Klik ikon Aplikasi. pengalih aplikasi, pilih Lakebase Postgres, lalu buat proyek baru dengan nama (misalnya, my-app-db) dan versi Postgres (terima Postgres 17 default). Untuk detail penyiapan lengkap, lihat Membuat proyek.

Tunggu hingga komputasi aktif (sekitar 1 menit) sebelum melanjutkan.

Langkah 2: Mengonfigurasi autentikasi dan skema database

Buat peran Postgres untuk perwakilan layanan aplikasi Anda dengan autentikasi OAuth, lalu buat tabel sampel dengan data untuk ditampilkan aplikasi Anda.

Menyiapkan autentikasi OAuth

Di proyek Lakebase Anda, buka Editor SQL dan jalankan perintah ini. Ekstensi ini databricks_auth memungkinkan autentikasi OAuth. Dengan itu, peran Postgres Anda menerima token Databricks alih-alih kata sandi tradisional:

-- Enable the Databricks authentication extension
CREATE EXTENSION IF NOT EXISTS databricks_auth;

-- Create a Postgres role for your app's service principal
-- Replace the UUID below with your DATABRICKS_CLIENT_ID from Step 1
SELECT databricks_create_role('<DATABRICKS_CLIENT_ID>', 'service_principal');

-- Grant necessary permissions (use the same DATABRICKS_CLIENT_ID)
GRANT CONNECT ON DATABASE databricks_postgres TO "<DATABRICKS_CLIENT_ID>";
GRANT CREATE, USAGE ON SCHEMA public TO "<DATABRICKS_CLIENT_ID>";

Ganti <DATABRICKS_CLIENT_ID> dengan nilai aplikasi DATABRICKS_CLIENT_ID Anda. Perwakilan layanan sekarang dapat mengautentikasi menggunakan token OAuth yang Azure Databricks kelola secara otomatis. Untuk detailnya, lihat Buat peran OAuth untuk identitas Azure Databricks.

Membuat skema database

Buat tabel sampel dengan izin eksplisit untuk perwakilan layanan Anda (perwakilan layanan tidak mewarisi izin skema default):

-- Create a sample table
CREATE TABLE notes (
    id SERIAL PRIMARY KEY,
    content TEXT NOT NULL,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);

-- Grant permissions to your app's service principal (use your DATABRICKS_CLIENT_ID)
GRANT SELECT, INSERT, UPDATE, DELETE ON TABLE notes TO "<DATABRICKS_CLIENT_ID>";

-- Insert sample data
INSERT INTO notes (content) VALUES
   ('Welcome to Lakebase Autoscaling!'),
   ('This app connects to Postgres'),
   ('Data fetched from your database');

Ganti <DATABRICKS_CLIENT_ID> dengan nilai Anda DATABRICKS_CLIENT_ID .

Langkah 3: Membangun dan mengonfigurasi aplikasi Anda

Unduh file aplikasi Anda, konfigurasikan koneksi database dengan rotasi token OAuth otomatis, dan uji secara lokal sebelum menyebarkan.

Mengunduh dan mengonfigurasi file aplikasi

Unduh file aplikasi Anda dari ruang kerja dengan menyalin perintah ekspor dari bagian Sinkronkan file aplikasi:

databricks workspace export-dir /Workspace/Users/<your-email>/databricks_apps/<app-folder>/flask-hello-world-app .

Edit app.yaml untuk menambahkan detail koneksi database. Dapatkan nilai koneksi Anda dari modal Lakebase Connect dengan memilih Parameter saja:

command: ['flask', '--app', 'app.py', 'run', '--host', '0.0.0.0', '--port', '8000']

env:
  - name: PGHOST
    value: '<your-endpoint-hostname>'
  - name: PGDATABASE
    value: 'databricks_postgres'
  - name: PGUSER
    value: '<DATABRICKS_CLIENT_ID>'
  - name: PGPORT
    value: '5432'
  - name: PGSSLMODE
    value: 'require'
  - name: ENDPOINT_NAME
    value: 'projects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id>'

Ganti tempat penampung:

  • <your-endpoint-hostname>: Salin nilai PGHOST dari modal Sambungkan (misalnya, ep-xyz.database.us-west-2.dev.databricks.com)
  • : Gunakan Anda dari Langkah 1
  • projects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id>: Di aplikasi Lakebase, buka tab Komputasi cabang Anda, klik Dapatkan ID untuk komputasi Anda, dan pilih Salin nama sumber daya.

Menerapkan rotasi token OAuth dan kueri database

Ganti app.py dengan kode ini, yang menambahkan rotasi token OAuth otomatis dan kueri database yang mengambil catatan dari Langkah 2:

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

app = Flask(__name__)

# Initialize Databricks client for token generation
w = WorkspaceClient()

# Custom connection class that generates fresh OAuth tokens
class OAuthConnection(psycopg.Connection):
    @classmethod
    def connect(cls, conninfo='', **kwargs):
        # Generate a fresh OAuth token for each connection (tokens are workspace-scoped)
        endpoint_name = os.environ["ENDPOINT_NAME"]
        credential = w.postgres.generate_database_credential(endpoint=endpoint_name)
        kwargs['password'] = credential.token
        return super().connect(conninfo, **kwargs)

# Configure connection parameters
username = os.environ["PGUSER"]
host = os.environ["PGHOST"]
port = os.environ.get("PGPORT", "5432")
database = os.environ["PGDATABASE"]
sslmode = os.environ.get("PGSSLMODE", "require")

# Create connection pool with automatic token rotation
pool = ConnectionPool(
    conninfo=f"dbname={database} user={username} host={host} port={port} sslmode={sslmode}",
    connection_class=OAuthConnection,
    min_size=1,
    max_size=10,
    open=True
)

@app.route('/')
def hello_world():
    # Use connection from pool (automatically gets fresh token)
    with pool.connection() as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT content, created_at FROM notes ORDER BY created_at DESC LIMIT 5")
            notes = cur.fetchall()

    # Display results
    notes_html = "<ul>" + "".join([f"<li>{note[0]} - {note[1]}</li>" for note in notes]) + "</ul>"
    return f'<h1>Hello from Lakebase!</h1><h2>Recent Notes:</h2>{notes_html}'

if __name__ == '__main__':
    app.run(host="0.0.0.0", port=8000)

Tiga komponen utama adalah:

  • WorkspaceClient: menghasilkan kredensial baru menggunakan SDK.
  • OAuthConnection: kelas koneksi kustom yang menyuntikkan kredensial baru ke setiap koneksi.
  • ConnectionPool: mengelola koneksi dan memanggil kelas kustom sesuai kebutuhan.

Untuk informasi selengkapnya tentang strategi rotasi kredensial dan penanganan kesalahan, lihat Contoh rotasi token.

Perbarui requirements.txt untuk menyertakan paket yang diperlukan:

flask
psycopg[binary,pool]
databricks-sdk>=0.81.0

Versi 0.81.0 atau lebih baru menyertakan metodegenerate_database_credential().

Menguji secara lokal

Uji aplikasi Anda secara lokal untuk memverifikasi bahwa koneksi database berfungsi sebelum menyebarkan. Saat menguji secara lokal, aplikasi berjalan sebagai akun pengguna Azure Databricks Anda (bukan perwakilan layanan), jadi ubah PGUSER ke alamat email Anda dalam variabel lingkungan di bawah ini.

Masuk ke ruang kerja Anda dan mengekspor variabel lingkungan.

databricks auth login

export PGHOST="<your-endpoint-hostname>"
export PGDATABASE="databricks_postgres"
export PGUSER="your.email@company.com"  # Use YOUR email for local testing, not the service principal
export PGPORT="5432"
export PGSSLMODE="require"
export ENDPOINT_NAME="<your-endpoint-name>"

Salin nilai dari app.yaml Anda, tetapi ganti nilai PGUSER (ID klien perwakilan layanan) dengan alamat email Azure Databricks Anda.

Instal dependensi dan jalankan aplikasi:

pip3 install --upgrade -r requirements.txt
python3 app.py

Buka http://localhost:8000 di browser Anda. Anda akan melihat "Halo dari Lakebase!" dengan tiga catatan sampel Anda. Kumpulan koneksi secara otomatis menghasilkan token OAuth baru saat membuat koneksi baru. Untuk detail selengkapnya, lihat Autentikasi token OAuth.

Output aplikasi lokal memperlihatkan

Langkah 4: Menyebarkan dan memverifikasi

Setelah menguji secara lokal, sinkronkan perubahan Anda ke folder ruang kerja dan sebarkan dari lokasi tersebut:

# Upload files to workspace
databricks sync . /Workspace/Users/<your-email>/my-lakebase-app

# Deploy from the uploaded location
databricks apps deploy <app-name> --source-code-path /Workspace/Users/<your-email>/my-lakebase-app

Ganti <your-email> dengan alamat email Azure Databricks Anda dan <app-name> dengan nama aplikasi Anda. Bendera --source-code-path memberi tahu penyebaran untuk menggunakan file yang Anda unggah alih-alih lokasi default aplikasi.

Tunggu hingga penyebaran selesai (2-3 menit), lalu akses aplikasi Anda di URL yang disediakan. Anda akan melihat "Halo dari Lakebase!" dengan catatan sampel Anda.

Lihat juga

  • Last updated on