Snorkel AI gives every pull request its own Lakebase copy of production

The problem: production-like tests were too slow to set up

Snorkel AI builds data development tools for AI teams, used by expert data annotators and frontier model providers. Its code ships through a CI pipeline built on Kubernetes, where every open pull request, plus ad-hoc commits, gets its own ephemeral namespace. Developers deploy their branch there and run the full test suite before merging.

Those tests are most useful against real data. Snorkel's production database is an Amazon Aurora PostgreSQL cluster, and restoring a copy of it for each PR namespace took about 20 minutes, according to Databricks. That is too slow when developers open PRs all day and the merge queue needs a fresh environment for every run. The alternatives, such as shared staging databases, trimmed fixtures or synthetic data, bring back the problems production-like testing is meant to avoid: tests interfering with each other, stale schemas, and bugs that only appear with real data shapes.

Snorkel set several success criteria: a new branch ready in under 60 seconds (p95), about 100 concurrent PR branches at peak, compatibility with the existing application over standard Postgres, a main branch within about 2 minutes of production, and private networking with Okta SSO.

How the branching setup works

Lakebase is Databricks' serverless Postgres database. Storage and compute are separate, and it supports branching: a new, isolated, writable branch can be created from any point in time, much like a Git branch. Branches are copy-on-write, so creating one is a metadata operation. A branch shares storage pages with its parent, and only the pages it changes are written as new storage.

Snorkel's pipeline has four steps:

• Keep main fresh. An AWS DMS task replicates the production Aurora database into the Lakebase main branch, with a full load first and then ongoing change data capture. This keeps main within about 2 minutes of production. Full LOB mode carries over JSONB values of 10 to 20 MB. The team already ran DMS, so no new replication pipeline was needed.
• Branch on PR open. The CI pipeline calls the Lakebase API to create a branch from main, sets an expiration time, and passes the connection details to the PR namespace.
• Test on real data. The application connects over standard Postgres and can run migrations, write data and change schemas without affecting production, main or any other PR.
• Clean up. The pipeline deletes the branch when the PR is merged or closed. If that step does not run, the branch expires when its TTL runs out.

Each branch has its own autoscaling compute that scales to zero between CI runs. Even with 100 provisioned branches, Snorkel expects only about 20 to 30 to serve traffic at once, at around 5 queries per second each at peak.

Security: private networking and separate identities

Two security requirements shaped the design. Production data moves from Aurora to Lakebase, and from Lakebase to the CI cluster, over private networking rather than the public internet.

For identity, Snorkel has two kinds of database users. Developers who debug a PR environment sign in with their usual Okta identity, which is federated to the Databricks account and provisioned through SCIM. CI pipelines run with no human involved, so each PR namespace runs as a Databricks service principal, scoped per team or per repo and also provisioned through SCIM. The PR pod uses the service principal's credentials to obtain short-lived Lakebase OAuth tokens through the Databricks SDK and refreshes them during long test runs. Databricks says there are no long-lived database passwords in CI. Postgres audit logs record the service principal identity for each query, so every query can be traced to the pipeline that ran it.

Cost, expiration and use by AI agents

Because branches share most data with main, storage grows with each branch's changes rather than with database size. Snorkel's PR branches write about 5 to 10 GB of changes on top of a ~1 TB base. Databricks says 100 concurrent PR databases on a ~2.5 TB production dataset add storage only for the data each branch changes.

Expiration matters because a branch keeps the version of data that existed when it was created. As DMS keeps updating main, a long-lived branch holds older pages and its storage grows. Setting a TTL on every CI branch removes it automatically once the PR is done, keeping storage close to the size of production plus the changes from active PRs.

Databricks argues the same pattern suits AI agents that write SQL, run migrations or change data. Each agent task could get its own branch of main, teams could review what an agent changed or simply delete the branch, and TTLs would clear abandoned ones. Agents could also run as service principals with short-lived OAuth tokens, so their queries can be traced in audit logs.

For teams with large production databases, the example shows how branching can make per-PR testing on real data practical without restores that take tens of minutes. The figures come from Databricks' account of Snorkel's own deployment.