Ir al contenido principal

Plugin de Lakebase

Plugin de Lakebase

Proporciona un pool de conexiones PostgreSQL para Lakebase Autoscaling de Databricks con refresh automático de tokens OAuth.

Características principales:

  • pg.Pool estándar, compatible con cualquier biblioteca PostgreSQL u ORM
  • Refresh automático de tokens OAuth (tokens de 1 hora, margen de refresh de 2 minutos)
  • Almacenamiento de tokens en caché para minimizar las llamadas a la API
  • Instrumentación integrada de OpenTelemetry (duración de las consultas, conexiones del pool, refresh de tokens)
  • Logger de AppKit configurado de forma predeterminada para los eventos de consulta y de conexión

Primeros pasos con Lakebase

La forma más sencilla de empezar a usar el plugin de Lakebase es utilizar la Databricks CLI para crear una nueva Databricks app con AppKit y el plugin de Lakebase instalados.

Requisitos previos

Pasos

  1. En primer lugar, crea un nuevo project de Lakebase Postgres con autoscaling siguiendo la documentación de introducción.
  2. Para agregar el plugin de Lakebase a tu project, ejecuta el comando databricks apps init y selecciona de forma interactiva el plugin Lakebase. La CLI te guiará para elegir un project, una branch y una base de datos de Lakebase.
    • Cuando se te solicite, selecciona Yes para desplegar la app en Databricks Apps justo después de crearla.

Uso básico

import { createApp, lakebase, server } from "@databricks/appkit";

await createApp({
  plugins: [server(), lakebase()],
});

Acceso al pool

Tras la inicialización, accede a Lakebase mediante el objeto AppKit.lakebase:

const AppKit = await createApp({
  plugins: [server(), lakebase()],
});

await AppKit.lakebase.query(`CREATE SCHEMA IF NOT EXISTS app`);

await AppKit.lakebase.query(`CREATE TABLE IF NOT EXISTS app.orders (
  id SERIAL PRIMARY KEY,
  user_id VARCHAR(255) NOT NULL,
  amount DECIMAL(10, 2) NOT NULL,
  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)`);

const result = await AppKit.lakebase.query(
  "SELECT * FROM app.orders WHERE user_id = $1",
  [userId],
);

// pg.Pool nativo (para ORMs o uso avanzado)
const pool = AppKit.lakebase.pool;

// Objetos de configuración listos para usar con ORMs
const ormConfig = AppKit.lakebase.getOrmConfig();  // { host, port, database, ... }
const pgConfig = AppKit.lakebase.getPgConfig();    // pg.PoolConfig

Configuración

Variables de entorno

Las variables de entorno requeridas son:

VariableDescripción
LAKEBASE_ENDPOINTRuta del recurso del endpoint (p. ej. projects/.../branches/.../endpoints/...)
PGHOSTHost de Lakebase (inyectado automáticamente en producción por el recurso postgres de Databricks Apps)
PGDATABASENombre de la base de datos (inyectado automáticamente en producción por el recurso postgres de Databricks Apps)
PGSSLMODEModo TLS: se debe establecer en require (inyectado automáticamente en producción por el recurso postgres de Databricks Apps)

Al desplegar en Databricks Apps con un recurso de base de datos postgres configurado, la plataforma inyecta automáticamente PGHOST, PGDATABASE, PGSSLMODE, PGUSER, PGPORT y PGAPPNAME. Solo es necesario establecer LAKEBASE_ENDPOINT de forma explícita:

env:
  - name: LAKEBASE_ENDPOINT
    valueFrom: postgres

Para el desarrollo local, el archivo .env se genera automáticamente mediante databricks apps init con los valores correctos de tu proyecto de Lakebase.

Para consultar la referencia completa de configuración (SSL, tamaño del pool, tiempos de espera, registro, ejemplos de ORM), consulta el README de @databricks/lakebase.

Configuración del pool

Pasa un objeto pool para sobrescribir cualquiera de los valores predeterminados:

await createApp({
  plugins: [
    lakebase({
      pool: {
        max: 10,                      // Máximo de conexiones del pool (valor por defecto: 10)
        connectionTimeoutMillis: 5000, // Tiempo de espera de conexión en ms (valor por defecto: 10000)
        idleTimeoutMillis: 30000,      // Tiempo de espera de conexión inactiva en ms (valor por defecto: 30000)
      },
    }),
  ],
});

On-Behalf-Of (OBO) — conexiones por usuario

Cuando tu aplicación necesita seguridad a nivel de fila (RLS) o aislamiento de datos por usuario, usa asUser(req) para ejecutar queries a través de un pool de conexiones de Lakebase propio de cada usuario. El pool de cada usuario se autentica con su identidad de Databricks, de modo que current_user de PostgreSQL refleja al usuario real.

Requisitos previos

  1. Habilita la autorización de usuario en tu Databricks App con el scope postgres. Consulta Autorización de usuario para ver las instrucciones de configuración. En tu databricks.yml:

    resources:
      apps:
        app:
          user_api_scopes:
            - postgres

    Las apps generadas con databricks apps init y el plugin de Lakebase ya lo incluyen automáticamente.

  2. Cada usuario de la app necesita un rol de Postgres en Lakebase. Crea uno con la Databricks CLI:

    databricks postgres create-role "projects/{project_id}/branches/{branch_id}" \
      --json '{"spec": {"identity_type": "USER", "postgres_role": "user@example.com"}}'

    También puedes crear los roles desde la interfaz de Lakebase, en Branch OverviewAdd role.

    note

    No otorgues databricks_superuser a los usuarios de OBO: los superusuarios omiten la RLS. Usa grants granulares en su lugar.

Uso

No se necesita configuración: basta con llamar a asUser(req):

const AppKit = await createApp({
  plugins: [server(), lakebase()],
});

// Consulta con el service principal (por defecto: omite RLS como propietario de la tabla)
const all = await AppKit.lakebase.query("SELECT * FROM app.orders");

// Consulta en el ámbito del usuario (pool por usuario, con RLS aplicado)
app.get("/api/my-orders", async (req, res) => {
  const result = await AppKit.lakebase
    .asUser(req)
    .query("SELECT * FROM app.orders ORDER BY created_at DESC");
  res.json(result.rows);
});

Cuando se llama a asUser(req):

  1. El token y la identidad del usuario se extraen de las cabeceras x-forwarded-access-token y x-forwarded-email (que Databricks Apps establece automáticamente).
  2. Se crea (o se reutiliza) un pg.Pool por usuario con las credenciales OAuth de ese usuario.
  3. query() y pool usan el pool del usuario: current_user en PostgreSQL refleja su identidad.

Ejemplo de seguridad a nivel de fila

-- Como service principal (durante la configuración de la app):
ALTER TABLE app.orders ENABLE ROW LEVEL SECURITY;

CREATE POLICY user_orders ON app.orders
  FOR ALL TO PUBLIC
  USING (owner = current_user);

-- Otorga acceso para que los usuarios OBO puedan consultar
GRANT USAGE ON SCHEMA app TO PUBLIC;
GRANT SELECT, INSERT ON ALL TABLES IN SCHEMA app TO PUBLIC;

Cómo funciona

  • El pool del service principal (AppKit.lakebase.pool) siempre se crea y se usa para operaciones DDL, carga inicial de datos y queries administrativas.
  • Los pools por usuario se crean en la primera llamada a asUser(req) y se almacenan en caché según la identidad del usuario. Cada pool tiene su propio ciclo de refresh del token OAuth.
  • Las conexiones inactivas dentro de los pools por usuario se cierran automáticamente (tiempo de inactividad de 30 s). Los objetos de pool vacíos se depuran de forma periódica.
  • Al apagar la aplicación, todos los pools (del service principal y de usuario) se cierran de forma ordenada.
  • En modo de desarrollo (NODE_ENV=development), si no hay ningún token de usuario disponible, asUser(req) recurre al pool del service principal y muestra una advertencia.
RLS y superusuarios

Los superusuarios de PostgreSQL omiten por completo la seguridad a nivel de fila (RLS). Los usuarios con el rol databricks_superuser verán todas las filas, sin importar las políticas de RLS. Para aplicar RLS, usa grant granular en lugar del rol de superusuario.

Permisos de base de datos

Cuando creas la app con el recurso de Lakebase siguiendo la guía de Primeros pasos, el service principal recibe automáticamente el permiso CONNECT_AND_CREATE sobre el recurso postgres. Esto le permite conectarse a la base de datos y crear objetos nuevos, pero no acceder a los esquemas o tablas existentes.

Desarrollo local

Para desarrollar en local contra una base de datos Lakebase desplegada:

  1. Despliega la app primero. El service principal crea el esquema y las tablas de la base de datos en el primer despliegue. Las apps generadas con databricks apps init se encargan de esto automáticamente: comprueban si las tablas existen al iniciar y omiten su creación si ya están.

  2. Concede databricks_superuser (omite este paso si eres el propietario del project de Lakebase, ya que tienes acceso completo):

    # Crear un nuevo rol con databricks_superuser
    databricks postgres create-role "projects/{project_id}/branches/{branch_id}" \
      --json '{"spec": {"identity_type": "USER", "postgres_role": "user@example.com", "membership_roles": ["DATABRICKS_SUPERUSER"]}}'

    Para conceder superusuario a un rol existente, usa update-role:

    databricks postgres update-role \
      "projects/{project_id}/branches/{branch_id}/roles/{role_id}" \
      "spec.membership_roles" \
      --json '{"spec": {"membership_roles": ["DATABRICKS_SUPERUSER"]}}'

    Como alternativa, puedes gestionar los roles en la interfaz de Lakebase Autoscaling, en la página Branch Overview de tu project → Add role / Edit role.

  3. Ejecuta en local: para la autenticación OAuth se usa tu identidad de usuario de Databricks (correo electrónico). El rol databricks_superuser otorga acceso DML completo (lectura/escritura de datos), pero no DDL (crear esquemas o tablas); por eso es importante desplegar primero (consulta la nota más abajo).

Para los demás usuarios, repite el paso 2 y crea un rol OAuth con databricks_superuser para cada uno.

tip

La autenticación por contraseña de Postgres es una alternativa más sencilla que evita la complejidad de los permisos de roles OAuth. Sin embargo, requiere que configures una contraseña para el usuario en la página Branch Overview de la interfaz de Lakebase Autoscaling.

¿Por qué desplegar primero?

Cuando se despliega la app, el service principal crea los esquemas y las tablas y pasa a ser su propietario. databricks_superuser otorga acceso DML completo (lectura/escritura) pero no DDL, por lo que el desarrollo local solo funciona una vez que el esquema existe.

Si ejecutas npm run dev primero, tus credenciales serán las propietarias del esquema y la app desplegada recibirá un error permission denied. Para solucionarlo, exporta primero los datos (pg_dump o una copia temporal del esquema), elimina el esquema y vuelve a desplegar. Tras el nuevo despliegue, el service principal recrea el esquema al iniciar. (En PostgreSQL, la propiedad de un esquema está ligada al rol que lo creó y los usuarios normales no pueden reasignarla).

Permisos granulares

Para la mayoría de los casos de uso, databricks_superuser es suficiente. Si en cambio necesitas grants a nivel de esquema, consulta la documentación oficial:

Script SQL para grants granulares

Despliega y ejecuta la app al menos una vez antes de aplicar estos grants, para que el service principal inicialice primero el esquema de la base de datos.

Reemplaza subject por el correo del usuario y schema por el nombre de tu esquema:

CREATE EXTENSION IF NOT EXISTS databricks_auth;

DO $$
DECLARE
  subject TEXT := 'your-subject';  -- Correo del usuario, como name@databricks.com
  schema TEXT := 'your_schema'; -- Reemplaza 'your_schema' por el nombre de tu esquema
BEGIN
  -- Crear rol OAuth para la identidad de Databricks
  PERFORM databricks_create_role(subject, 'USER');

  -- Acceso a la conexión y al esquema
  EXECUTE format('GRANT CONNECT ON DATABASE "databricks_postgres" TO %I', subject);
  EXECUTE format('GRANT ALL ON SCHEMA %s TO %I', schema, subject);

  -- Privilegios sobre objetos existentes
  EXECUTE format('GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA %s TO %I', schema, subject);
  EXECUTE format('GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA %s TO %I', schema, subject);
  EXECUTE format('GRANT ALL PRIVILEGES ON ALL FUNCTIONS IN SCHEMA %s TO %I', schema, subject);
  EXECUTE format('GRANT ALL PRIVILEGES ON ALL PROCEDURES IN SCHEMA %s TO %I', schema, subject);

  -- Privilegios predeterminados sobre objetos futuros
  EXECUTE format('ALTER DEFAULT PRIVILEGES IN SCHEMA %s GRANT ALL ON TABLES TO %I', schema, subject);
  EXECUTE format('ALTER DEFAULT PRIVILEGES IN SCHEMA %s GRANT ALL ON SEQUENCES TO %I', schema, subject);
  EXECUTE format('ALTER DEFAULT PRIVILEGES IN SCHEMA %s GRANT ALL ON FUNCTIONS TO %I', schema, subject);
  EXECUTE format('ALTER DEFAULT PRIVILEGES IN SCHEMA %s GRANT ALL ON ROUTINES TO %I', schema, subject);
END $$;

Databricks Developer Hub

¿Todo listo para lanzar tu próxima aplicación basada en agentes en minutos?

Leer la documentación