Skip to navigation

Databricks Lakebase Connector

View as Markdown

Lakebase is Databricks’ managed PostgreSQL. Reach it the same way as any PostgreSQL source — through PostgresSource — but authenticate with sdm_connectors.databricks.lakebase(), which mints a short-lived token as your workspace identity instead of a static password.

Installation

Install the lakebase extra, not postgres:

pip install "structured-data-models-client[relational,lakebase]"

On Databricks aarch64 compute, the postgres extra’s psycopg[binary] bundles an OpenSSL build that fails its FIPS self-test and aborts the process. The lakebase extra brings plain psycopg, which uses the runtime’s own libpq and works on both architectures. For the same reason, do not install [all] on Databricks aarch64 compute — it pulls in postgres and, with it, psycopg[binary].

Sample in place with PostgresSource

from databricks.sdk import WorkspaceClient
from sdm_client import PredictiveQuery, Schema
from sdm_connectors import PostgresSource
from sdm_connectors.databricks import lakebase
workspace = WorkspaceClient()
connection = lakebase(endpoint, database="databricks_postgres", workspace=workspace)
source = PostgresSource(connection, schema="public")
schema = Schema.from_source(source, {"users": "users", "orders": "orders"})
request = PredictiveQuery(
"PREDICT COUNT(orders.*, 0, 30, days) FOR EACH users.user_id", schema
).materialize(indices=[7, 12, 3])

endpoint is the Lakebase endpoint’s full name (projects/.../branches/.../endpoints/...). lakebase() mints the password as a database credential for that endpoint and connects to the host it names, unless you pass host= explicitly; workspace defaults to an ambient WorkspaceClient(). Defaults are sslmode='require' and autocommit=True, both overridable through extra keyword arguments. The connection is yours — lakebase() does not close it.

Reach Lakebase from Databricks serverless compute, which connects on a private address. A classic cluster connects over the public internet and only succeeds if the workspace admits its egress address.

Every connection lakebase() opens carries nvidia_structured-data-models-client/<version> as its application_name, as Databricks requires of a partner product.

Review the inferred schema and relationships before predicting — see Build a Schema.