Migrate from PostgreSQL to Databricks Lakebase

Create the project

Check the source version with psql "$SRC_URL" -c 'show server_version' and create the project on the same major version:

export SRC_URL="postgresql://aidbox:<password>@localhost:5432/aidbox"
export DATABRICKS_WS="https://<workspace>.cloud.databricks.com"
export PROJECT_ID="aidbox"

databricks postgres create-project "$PROJECT_ID" \
  --json '{"spec": {"display_name": "Aidbox", "pg_version": "17"}}'

# A new project has the "production" branch with the read-write endpoint "primary"
export BRANCH="projects/$PROJECT_ID/branches/production"
export ENDPOINT="$BRANCH/endpoints/primary"
export LB_HOST=$(databricks postgres get-endpoint "$ENDPOINT" -o json | jq -r .status.hosts.host)

LB_HOST looks like ep-<name>.database.<region>.cloud.databricks.com.

Disable scale-to-zero and size the compute

A suspended compute closes the Aidbox connection pool. Pick the CU range from the load on the source database. The spread between minimum and maximum can't exceed 16 CU.

databricks postgres update-endpoint "$ENDPOINT" spec.suspension \
  --json '{"spec": {"no_suspension": true}}'

databricks postgres update-endpoint "$ENDPOINT" \
  "spec.autoscaling_limit_min_cu,spec.autoscaling_limit_max_cu" \
  --json '{"spec": {"autoscaling_limit_min_cu": 1.0, "autoscaling_limit_max_cu": 8.0}}'

Create a service principal

databricks service-principals create --display-name aidbox-lakebase -o json
# Note "id" and "applicationId" from the output
databricks service-principal-secrets-proxy create <id> -o json
# Store "secret" in your secret manager; Databricks shows it once
export SP_CLIENT_ID=<applicationId>
export SP_CLIENT_SECRET=<secret>

Create a Postgres role for the service principal

The role name is the application ID.

databricks postgres create-role "$BRANCH" --role-id aidbox-sp \
  --json "{\"spec\": {\"identity_type\": \"SERVICE_PRINCIPAL\", \"postgres_role\": \"$SP_CLIENT_ID\"}}"

Grant the role access to the database

Connect as the project owner with a short-lived token. Aidbox creates and alters its own tables, so the role needs CREATE on the database and on the public schema.

export PGPASSWORD=$(databricks postgres generate-database-credential "$ENDPOINT" -o json | jq -r .token)

psql "postgresql://<your-email>@$LB_HOST:5432/databricks_postgres?sslmode=require" <<SQL
GRANT CONNECT, CREATE, TEMPORARY ON DATABASE databricks_postgres TO "$SP_CLIENT_ID";
GRANT ALL ON SCHEMA public TO "$SP_CLIENT_ID";
SQL

URL-encode the @ in your email as %40.

Create extensions as the service principal

The Databricks CLI authenticates as the service principal when you set DATABRICKS_HOST, DATABRICKS_CLIENT_ID and DATABRICKS_CLIENT_SECRET. Create every extension the source database has that Lakebase supports, so the service principal owns them.

sp_token() {
  DATABRICKS_HOST=$DATABRICKS_WS DATABRICKS_CLIENT_ID=$SP_CLIENT_ID DATABRICKS_CLIENT_SECRET=$SP_CLIENT_SECRET \
    databricks postgres generate-database-credential "$ENDPOINT" -o json | jq -r .token
}
export LB_SP_URL="postgresql://$SP_CLIENT_ID@$LB_HOST:5432/databricks_postgres?sslmode=require"
export PGPASSWORD=$(sp_token)

psql "$LB_SP_URL" <<'SQL'
CREATE EXTENSION IF NOT EXISTS pg_trgm            SCHEMA public;
CREATE EXTENSION IF NOT EXISTS unaccent           SCHEMA public;
CREATE EXTENSION IF NOT EXISTS pgcrypto           SCHEMA public;
CREATE EXTENSION IF NOT EXISTS fuzzystrmatch      SCHEMA public;
CREATE EXTENSION IF NOT EXISTS pg_stat_statements SCHEMA public;
CREATE EXTENSION IF NOT EXISTS postgis            SCHEMA public; -- only if the source has it
SQL

Check: psql "$LB_SP_URL" -c 'select current_user' returns the application ID, and \dx lists the extensions.

Stop Aidbox and dump the database

docker compose stop aidbox

pg_dump -Fc -v -d "$SRC_URL" -f aidbox.dump

Build the restore list

The target already has the extensions. spatial_ref_sys belongs to PostGIS on the target, and the service principal can't write to it.

pg_restore -l aidbox.dump \
  | grep -v -E ' EXTENSION |TABLE DATA public spatial_ref_sys' > restore.list

Restore as the service principal

--no-owner --no-acl makes the service principal the owner of every restored object. Each -j worker opens its own connection with the token from PGPASSWORD, and Lakebase checks the token at login, so generate a fresh one right before the restore.

export PGPASSWORD=$(sp_token)

pg_restore -v --no-owner --no-acl -j 4 -L restore.list \
  -d "$LB_SP_URL" aidbox.dump 2> restore.log

grep -E '^pg_restore: error' restore.log
psql "$LB_SP_URL" -c 'ANALYZE'

ANALYZE prints permission denied to analyze "pg_..." warnings for system catalogs. You can ignore them.

Compare row counts

psql "$LB_SP_URL" -At -c "select format('select %L, count(*) from %I.%I;', table_name, table_schema, table_name)
  from information_schema.tables where table_schema = 'public' and table_type = 'BASE TABLE' order by 1" \
  | psql "$LB_SP_URL" -At > counts_target.txt

diff counts_source.txt counts_target.txt && echo "counts match"