Skip to main content

Managing Databricks Lakebase Schemas with Atlas

Databricks Lakebase is a managed PostgreSQL service in a Databricks workspace. A Lakebase project holds branches, and each branch serves its Postgres databases through a compute endpoint. Atlas connects to an endpoint over postgres://, the same as any other PostgreSQL database.

Lakebase accepts two kinds of login:

  • An OAuth token, issued to a Databricks user or service principal. A token expires one hour after it is issued. The databricks_lakebase_token data source generates a new one each time Atlas loads atlas.hcl.
  • A native Postgres password, for a role created with one. Passwords do not expire.

This guide covers both, and the objects Lakebase creates in every database that the declarative workflow has to leave alone.

Prerequisites​

  1. A Lakebase project, and a Postgres role on its branch for the identity Atlas connects as. Lakebase creates an OAuth role for the project owner. For any other user or service principal, create an OAuth role and grant it privileges on the schemas Atlas manages, as described in Manage Postgres roles.
  2. The Databricks CLI and jq, to look up the endpoint.
  3. Docker, for the dev database.
  4. Atlas installed on your machine:

To download and install the latest release of the Atlas CLI, simply run the following in your terminal:

curl -sSf https://atlasgo.sh | sh
Atlas Pro Feature

The databricks_lakebase_token data source is available to Atlas Pro users. Run the following command to use it:

atlas login

Finding the Endpoint​

The data source takes the endpoint path, in the form projects/{project}/branches/{branch}/endpoints/{endpoint}. The connection URL takes the endpoint host. databricks postgres list-endpoints prints both for every endpoint on a branch. It takes the branch path, without the /endpoints/{endpoint} part:

databricks postgres list-endpoints projects/my-app/branches/production -o json \
  | jq -r '.[] | "\(.name) \(.status.hosts.host)"'

Each output line holds an endpoint path, followed by its host. If you already have the endpoint path, databricks postgres get-endpoint prints its host:

databricks postgres get-endpoint projects/my-app/branches/production/endpoints/primary -o json \
  | jq -r .status.hosts.host

Use this host in the URL, not the -pooler host that the endpoint also lists. The pooler rejects OAuth tokens.

The dev database runs the same Postgres version as the project. Print the version:

databricks postgres get-project projects/my-app -o json | jq .status.pg_version
17

Connecting with an OAuth Token​

The lakebase env below logs in with a token from the databricks_lakebase_token data source:

atlas.hcl
locals {
  host = "ep-example.database.us-west-2.cloud.databricks.com"
  user = "user@example.com"
}

data "databricks_lakebase_token" "db" {
  endpoint = "projects/my-app/branches/production/endpoints/primary"
}

env "lakebase" {
  url = urluserinfo(
    "postgres://${local.host}:5432/databricks_postgres?sslmode=require&search_path=public",
    local.user,
    data.databricks_lakebase_token.db,
  )
  dev = "docker://postgres/17/dev?search_path=public"
  // Functions that Lakebase manages in "public". See "Objects Lakebase Manages".
  exclude = [
    "grant_usage_on_new_schema[type=function]",
    "grant_all_on_new_sequences[type=function]",
    "grant_select_on_new_objects[type=function]",
  ]
  schema {
    src = "file://schema.sql"
  }
}
  • endpoint is the endpoint path from Finding the Endpoint.
  • local.user is the Postgres role of the identity the token is issued to. For a Databricks user, the role is the user's email. For a service principal, it is the application ID. A token presented with any other role fails to log in.
  • urluserinfo sets the user and password of the URL, and escapes them. A Databricks email contains @, which must be escaped in a URL.
  • databricks_postgres is the database that Lakebase creates on every branch.
  • search_path=public binds the env to the public schema. The dev URL is bound to public too, so both URLs have the same scope.

Inspect the schema to check the connection:

atlas schema inspect --env lakebase

Atlas generates a token each time it loads the config, so each command starts with a token that has not expired.

Choosing Databricks credentials​

The data source requests the token from Databricks, so it needs Databricks credentials for the identity. The user in the URL is only the Postgres role, and does not sign in to Databricks. The data source finds the credentials with unified authentication. Both attributes that select them are optional:

  • profile names the ~/.databrickscfg profile. It defaults to DATABRICKS_CONFIG_PROFILE, or the DEFAULT profile.
  • host sets the workspace URL. It defaults to DATABRICKS_HOST, or the host of the profile.

Sign in once in the browser. The CLI saves a profile and refreshes its token:

databricks auth login --host https://dbc-a1b2c3d4-e5f6.cloud.databricks.com --profile lakebase

Point the data source at the profile, or set DATABRICKS_CONFIG_PROFILE=lakebase instead:

atlas.hcl
data "databricks_lakebase_token" "db" {
  endpoint = "projects/my-app/branches/production/endpoints/primary"
  profile  = "lakebase"
}

Atlas does not run the browser login itself. A profile created by databricks auth login gets its token from the Databricks CLI, so the databricks binary must be on the PATH wherever Atlas runs. The connection user is your Databricks email.

Connecting with a Password​

A native password role logs in without Databricks credentials. Lakebase disables password connections in new projects. Enable them as described in Manage roles. Then create a password role, as described in Manage Postgres roles. Lakebase shows the password once, so store it in your secret store when you create the role.

Grant the role the same privileges an OAuth role needs on the schemas Atlas manages. Then read the role and the password from the environment:

atlas.hcl
locals {
  host = "ep-example.database.us-west-2.cloud.databricks.com"
}

env "lakebase" {
  url = urluserinfo(
    "postgres://${local.host}:5432/databricks_postgres?sslmode=require&search_path=public",
    getenv("LAKEBASE_USER"),
    getenv("LAKEBASE_PASSWORD"),
  )
  dev = "docker://postgres/17/dev?search_path=public"
  exclude = [
    "grant_usage_on_new_schema[type=function]",
    "grant_all_on_new_sequences[type=function]",
    "grant_select_on_new_objects[type=function]",
  ]
  schema {
    src = "file://schema.sql"
  }
}

A password stays valid until you rotate it. Prefer OAuth tokens where the identity can use them.

Objects Lakebase Manages​

Lakebase creates these objects in every database on a branch:

ObjectTypeSchema
__db_systemSchema
grant_usage_on_new_schema, grant_all_on_new_sequences, grant_select_on_new_objectsFunctionpublic
on_create_schema, on_create_sequence, on_create_table_or_viewEvent trigger

The cloud_admin system role owns the functions and the event triggers. The event triggers call the functions, which grant Databricks system roles access to each new schema, table, view, and sequence. Atlas does not report these grants as schema changes.

The schema file is authoritative for the schemas in scope. Without the exclude line, every schema apply on public plans to drop the three functions. This plan ran against a database that already has the users table, and is abridged to the planned statements:

Planning migration statements (3 in total):

  -- drop "grant_all_on_new_sequences" function:
    -> DROP FUNCTION "grant_all_on_new_sequences";
  -- drop "grant_select_on_new_objects" function:
    -> DROP FUNCTION "grant_select_on_new_objects";
  -- drop "grant_usage_on_new_schema" function:
    -> DROP FUNCTION "grant_usage_on_new_schema";

The event triggers depend on the functions, so Postgres rejects the first DROP FUNCTION, and the apply fails:

Error: executing statement "DROP FUNCTION \"grant_all_on_new_sequences\";": pq: cannot drop function grant_all_on_new_sequences() because other objects depend on it (2BP01)

Exclude the three functions by name to leave them out of the plan:

atlas.hcl
env "lakebase" {
  // ...
  exclude = [
    "grant_usage_on_new_schema[type=function]",
    "grant_all_on_new_sequences[type=function]",
    "grant_select_on_new_objects[type=function]",
  ]
}

List the names rather than a pattern such as grant_*. A pattern also matches your own functions that start with grant_, and Atlas leaves them out of every plan without a warning.

Each name is matched within public, because search_path=public binds the URL to it. A URL without search_path covers the whole database, including __db_system and the event triggers. Bind the env to the schemas you manage.

The databricks_auth extension, which creates OAuth roles in SQL, also installs functions in public. These functions belong to the extension, so Atlas leaves them out of the plan without an exclude.

The versioned workflow runs only the statements in the migration files, so it does not need the exclude.

Applying Changes​

Both workflows read the desired schema from schema.sql:

schema.sql
CREATE TABLE users (
  id bigint PRIMARY KEY GENERATED ALWAYS AS IDENTITY,
  created_at timestamptz NOT NULL DEFAULT now(),
  email text NOT NULL UNIQUE
);

atlas schema apply computes the plan against the live database:

atlas schema apply --env lakebase

Atlas prints the plan and waits for approval before it applies anything. On an empty public schema, the plan creates the table. The output is abridged to the planned statements:

Planning migration statements (1 in total):

  -- create "users" table:
    -> CREATE TABLE "users" (
         "id" bigint NOT NULL GENERATED ALWAYS AS IDENTITY,
         "created_at" timestamptz NOT NULL DEFAULT now(),
         "email" text NOT NULL,
         PRIMARY KEY ("id"),
         CONSTRAINT "users_email_key" UNIQUE ("email")
       );

Add --dry-run to print the plan without applying it.

Troubleshooting​

ErrorCauseFix
default auth: cannot configure default credentialsAtlas found no Databricks credentials.Run databricks auth login and set profile, or set the service principal variables.
default auth: databricks CLI not foundThe profile gets its token from the Databricks CLI, and the CLI is not on the PATH.Install the CLI where Atlas runs.
Endpoint name expects 'projects/{project_id}/branches/{branch_id}/endpoints/{endpoint_id}' formatendpoint is not a full endpoint path.Use the path from Finding the Endpoint.
OAuth authentication is not supported with pooled connectionsThe URL uses the -pooler host.Use the host without -pooler.
cannot drop function grant_... because other objects depend on it (2BP01)schema apply planned to drop a function that Lakebase manages.Exclude the three functions by name, as shown in Objects Lakebase Manages.

Next Steps​

Secrets and IAM Authentication

Read database credentials from secret stores and generate cloud IAM tokens in atlas.hcl.

Database URLs

Schema scope and database scope in connection URLs, and how they change the planned DDL.

Declarative Migrations

Plan and apply schema changes from the desired state with atlas schema apply.

Versioned Migrations

The migration directory, planning with migrate diff, and applying with migrate apply.

GitHub logoGitHub Actions

Atlas Actions for linting migrations on a pull request and applying them on merge.

Databricks Unity Catalog

Manage Unity Catalog schemas in a Databricks workspace with the Databricks driver.