Подключение пользовательского приложения Databricks к Lakebase

В этом руководстве показано, как подключить приложение Databricks к Lakebase Autoscaling с автоматическим поворотом учетных данных. Приложение создает новые учетные данные базы данных из Databricks до истечения срока их действия. В примере используется Flask, но шаблон проверки подлинности применяется к любой платформе.

Принцип работы

Databricks Apps аутентифицируются в Lakebase с помощью маркеров OAuth, срок действия которых истекает через час. Для этого необходимо создать роль Postgres для субъекта-службы приложения, а затем настроить приложение для автоматического создания новых маркеров при необходимости подключения к базе данных. Приложение использует пул подключений. Пул создает новые подключения с новыми маркерами по мере необходимости, поэтому приложение никогда не использует учетные данные с истекшим сроком действия.

При развертывании приложения в Azure Databricks оно запускается в качестве учетной записи службы и создает маркеры для этого удостоверения. При локальном тестировании приложение запускается в качестве учетной записи пользователя Azure Databricks и создает маркеры для вас. Оба используют один и тот же код поворота маркеров. Изменяется только контекст аутентификации.

Перед началом работы

Чтобы завершить работу с этим руководством, вам потребуется:

Шаг 1. Создание приложения и базы данных

Сначала создайте приложение Databricks и проект Lakebase. Приложение автоматически получает учетную запись служебного принципала, чтобы использовать ее для проверки подлинности базы данных.

Создание приложения

Создайте новое приложение Databricks с помощью шаблона Flask Hello world. См. статью "Создание приложения Databricks" из шаблона.

После создания приложения перейдите на вкладку "Среда приложения" и запишите DATABRICKS_CLIENT_ID значение (формат UUID, например 6b215d2b-f099-4bdb-900a-60837201ecec). Это становится именем пользователя Postgres приложения для проверки подлинности OAuth.

Замечание

Пока не развертывайте приложение. Сначала настройте подключение к базе данных.

Создание базы данных

Создайте новый проект для хостинга вашей базы данных с использованием Lakebase Autoscaling. Щелкните значок приложения в переключателе приложений, выберите Lakebase Postgres, а затем создайте новый проект с именем (например, my-app-db) и версией Postgres (примите значение по умолчанию, Postgres 17). Полные сведения о настройке см. в разделе "Создание проекта".

Дождитесь, пока вычислительные ресурсы станут активными (примерно через 1 минуту), прежде чем продолжить.

Шаг 2. Настройка проверки подлинности и схемы базы данных

Создайте роль Postgres для учетной записи службы вашего приложения с аутентификацией OAuth, а затем создайте пример таблицы с данными для отображения вашим приложением.

Настройка проверки подлинности OAuth

В проекте Lakebase откройте редактор SQL и выполните эти команды. Расширение databricks_auth включает проверку подлинности OAuth. Благодаря этому ваши роли Postgres принимают токены Databricks вместо традиционных паролей.

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

Замените <DATABRICKS_CLIENT_ID> значением приложения DATABRICKS_CLIENT_ID . Теперь служебный принципал может аутентифицироваться с помощью маркеров OAuth, которые Azure Databricks автоматически управляет. Дополнительные сведения см. в разделе Создание роли OAuth для идентификации Azure Databricks.

Создание схемы базы данных

Создайте пример таблицы с явными разрешениями для субъекта-службы (субъекты-службы не наследуют разрешения схемы по умолчанию):

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

Замените <DATABRICKS_CLIENT_ID> значением DATABRICKS_CLIENT_ID .

Шаг 3. Создание и настройка приложения

Скачайте файлы приложения, настройте подключение к базе данных с автоматическим поворотом маркера OAuth и протестируйте локально перед развертыванием.

Скачивание и настройка файлов приложений

Скачайте файлы приложения из рабочей области, скопируйте команду экспорта из раздела "Синхронизация файлов " приложения:

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

Измените app.yaml, чтобы добавить сведения о подключении к базе данных. Получите значения подключения из модала Lakebase Connect , выбрав только параметры:

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>'

Замените заполнители:

  • <your-endpoint-hostname>: скопируйте значение PGHOST из модального модуля Connect (например, ep-xyz.database.us-west-2.dev.databricks.com)
  • <DATABRICKS_CLIENT_ID>: Используйте DATABRICKS_CLIENT_ID из Шага 1
  • projects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id>: в приложении Lakebase перейдите на вкладку "Компьютеры" вашей ветви, нажмите кнопку "Получить идентификатор" для вашего компьютера и выберите "Копировать имя ресурса".

Ротирование OAuth токена и выполнение запросов к базе данных.

Замените app.py этим кодом, который добавляет автоматическую ротацию токена OAuth и запрос к базе данных, извлекающий заметки из Step 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)

Ниже перечислены три основных компонента:

  • WorkspaceClient: создает новые учетные данные с помощью пакета SDK.
  • OAuthConnection: пользовательский класс подключения, который внедряет новые учетные данные в каждое подключение.
  • ConnectionPool: управляет подключениями и вызывает пользовательский класс по мере необходимости.

Дополнительные сведения о стратегиях смены учетных данных и обработке ошибок см. в примерах вращения токенов.

Обновите requirements.txt, чтобы включить необходимые пакеты.

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

Версия 0.81.0 или более поздней включает generate_database_credential() метод.

Локальное тестирование.

Протестируйте приложение локально, чтобы убедиться, что подключение к базе данных работает перед развертыванием. При локальном тестировании приложение выполняется как пользовательская учетная запись Azure Databricks (а не учетная запись службы), поэтому измените PGUSER на адрес вашей электронной почты в переменных среды ниже.

Выполните проверку подлинности в рабочей области и экспортируйте переменные среды.

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

Скопируйте значения из app.yaml, но замените значение PGUSER (идентификатор клиента сервисного принципала) на адрес электронной почты Azure Databricks.

Установите зависимости и запустите приложение:

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

Откройте http://localhost:8000 в браузере. Вы должны увидеть "Привет из Lakebase!" с вашими тремя примерами заметок. Пул подключений автоматически создает новые маркеры OAuth при создании новых подключений. Дополнительные сведения см. в разделе проверки подлинности маркера OAuth.

Вывод локального приложения, показывающий

Шаг 4. Развертывание и проверка

После локального тестирования синхронизируйте изменения в папке рабочей области и развертывайте из этого расположения:

# 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

Замените <your-email> адресом электронной почты Azure Databricks и <app-name> именем приложения. Флаг --source-code-path указывает развертыванию использовать загруженные вами файлы вместо расположения приложения по умолчанию.

Дождитесь завершения развертывания (2–3 минуты), а затем перейдите к приложению по указанному URL-адресу. Вы увидите "Hello from Lakebase!" с примерами заметок.

См. также