Databricks Schema Management with Atlas Demo

This repository contains the code shown in our Databricks Schema Management with Atlas demo video. To set up with your own Databricks workspace, follow the instructions below.

For more about using Databricks with Atlas, check out our Databricks guide.

What's inside?

• Setting up
• Local declarative workflow
• Local versioned workflow
• Using Atlas in CI/CD
  • Setting up
  • CI/CD for declarative migrations
  • CI/CD for versioned migrations

Prerequisites

1. Databricks workspace and SQL warehouse:

• A Databricks workspace
• A running SQL warehouse or compute cluster
• A Personal Access Token (PAT) for authentication

2. Atlas installed on your machine:

curl -sSf https://atlasgo.sh | ATLAS_VERSION="beta" sh

3. An Atlas Pro account. To create or log in to an account, run:

atlas login

Setting up Databricks authentication

Creating a Personal Access Token

1. Go to your Databricks workspace
2. Click on your username in the top-right corner and select "User Settings"
3. Go to the "Developer" tab
4. Click "Manage" next to "Access tokens"
5. Click "Generate new token"
6. Give your token a description and set an expiration (optional)
7. Copy the generated token and store it securely

See the Databricks documentation for more details.

Environment setup

To get DATABRICKS_HOST and DATABRICKS_WAREHOUSE, refer to the Databricks documentation.

Set your Databricks credentials as environment variables:

export DATABRICKS_TOKEN="your-personal-access-token"
export DATABRICKS_HOST="dbc-xxxxxxxx-xxxx.cloud.databricks.com"
export DATABRICKS_WAREHOUSE="/sql/1.0/warehouses/your-warehouse-id"
export DATABASE_NAME="name-of-your-target-db" # If not the default catalog

Use these credentials to create your connection URL:

export DATABRICKS_URL="databricks://$DATABRICKS_TOKEN@$DATABRICKS_HOST:443$DATABRICKS_WAREHOUSE?catalog=$DATABASE_NAME"

Create a dev database

Atlas's automatic migrations flow uses a secondary dev database to simulate the schema changes.

Create a new catalog in your Databricks workspace to use as a dev database:

CREATE CATALOG IF NOT EXISTS demo_dev;

Make another connection URL for the dev database:

export DATABRICKS_DEV_URL="databricks://$DATABRICKS_TOKEN@$DATABRICKS_HOST:443$DATABRICKS_WAREHOUSE?catalog=demo_dev"

Configuration file (optional)

This repository includes an atlas.hcl configuration file containing an environment called demo. It consists of the variables needed to run most Atlas commands, so you can add the flag --env demo to your commands instead of each individual variable flag (e.g., -u, --dev-url, etc.).

The guided steps below will include each variable flag in the commands, but you may use the --env flag if you prefer. The declarative workflow CI/CD section will use this environment in the workflow file rather than writing out each variable.

Inspecting your target database

To begin defining our desired state as code, we need to retrieve the current schema definition to use as the base for our changes. We call this inspecting our database.

To inspect our Databricks schema and store it in SQL format, run:

atlas schema inspect -u "$DATABRICKS_URL" --format '{{ sql . }}' > schema.sql

schema.sql is the file we will alter to define the desired state we want to migrate our target database to.

Declarative migrations

Atlas supports a workflow called declarative schema migrations. In this workflow, you first define the desired state of your database schema. Then, Atlas calculates the diff between the desired state and the actual state of your database, and generates the SQL commands that will bring your database to the desired state.

To use declarative migrations on your Databricks catalog:

1. Push your desired schema to Atlas Cloud by running:

atlas schema push databricks-demo -u "file://schema.sql" --dev-url "$DATABRICKS_DEV_URL"

Replace databricks-demo with your desired Atlas Cloud repository name.

2. Update your desired schema definition by adding, removing, and/or altering the SQL statements.
3. Run the following command to create a migration plan for approval:

atlas schema apply -u "$DATABRICKS_URL" \
  --to file://schema.sql \
  --tx-mode="none" \
  --dev-url "$DATABRICKS_DEV_URL"

4. Select Approve and apply to apply the migration to your target catalog.
5. Run the atlas schema push command again to update the Atlas Cloud repository.

Versioned migrations

Atlas also supports a workflow called versioned schema migrations. In this workflow, the "current state" of your database is defined by a directory of migration files. When you make changes to the desired state of your database schema, Atlas writes a new file to the migrations directory that takes the "current state" to the desired state.

Head to our docs to read more about Declarative vs Versioned Migrations.

Important

For this demo, we are beginning with an empty database. If your database and schema.sql file are already populated with objects beyond the default schema, start by creating a baseline migration, then proceed to Step 2.

To use versioned migrations on your Databricks catalog:

1. Create a migrations directory by running the following command:

atlas migrate new init --edit

Exit by typing :q and hitting Enter.

2. Push the migration directory to Atlas Cloud by running:

atlas migrate push databricks-demo --dev-url "$DATABRICKS_DEV_URL"

Replace databricks-demo with your desired Atlas Cloud repository name.

3. Update your desired schema definition by adding, removing, and/or altering the SQL statements.
4. Run the following command to create a migration file:

atlas migrate diff --to file://schema.sql --dev-url "$DATABRICKS_DEV_URL"

5. Run the following command to apply the migration file to your target catalog:

atlas migrate apply -u "$DATABRICKS_URL" --tx-mode="none"

6. Run the atlas migrate push command again to update the Atlas Cloud repository.

Note

The migrations directory is created in the working directory and named migrations by default. If you'd like to use a different directory for the migration files, add --dir "file://<path/to/dir>" to all of the above commands.

Automating migrations with GitHub Actions

Atlas offers AtlasAction tasks that can be used to automate the migration creation process based on our database CI/CD principles just by opening and merging pull requests on GitHub.

Prerequisites

1. A GitHub repository containing your Atlas code setup and the Actions workflow file (in following steps)
2. An Atlas Cloud bot token stored as a GitHub repository secret
3. The connection URLs stored as GitHub repository secrets

Create an Atlas Cloud bot token

Atlas Cloud uses a bot token to connect with your CI/CD pipeline for reporting CI runs to Atlas Cloud and pushing updated schemas or migration directories to Atlas Cloud repositories.

Follow these instructions to create a bot token.

Then, store it as a secret in your GitHub repository by running:

gh secret set ATLAS_CLOUD_TOKEN

Store connection URLs as secrets

Get your connection URLs as strings by running printenv for both DATABRICKS_URL and DATABRICKS_DEV_URL.

For example:

printenv DATABRICKS_URL

databricks://your-personal-access-token@dbc-xxxxxxxx-xxxx.cloud.databricks.com:443/sql/1.0/warehouses/your-warehouse-id?catalog=name-of-your-target-db

Store them as secrets by running

gh secret set DATABRICKS_URL

and

gh secret set DATABRICKS_DEV_URL

Declarative workflow in CI/CD

First, set up an Atlas Cloud repository by running:

atlas schema push databricks-cicd-demo --dev-url "$DATABRICKS_DEV_URL" -u "file://schema.sql"

Replace databricks-demo with your Atlas Cloud repository name.

When you update the desired schema file, push it to a new branch in your GitHub repository and make a pull request. This will trigger the workflow file.

Workflow file

Copy the .github/workflows/ci-atlas-declarative.yaml file into your repository, replacing databricks-demo with your Atlas Cloud repository name.

Let's go over the jobs.

plan:
  if: ${{ github.event_name == 'pull_request' }}
  runs-on: ubuntu-latest
  env:
    GITHUB_TOKEN: ${{ github.token }}
  steps:
    - uses: actions/checkout@v3
      with:
        fetch-depth: 0
    - uses: ariga/setup-atlas@v0
      with:
        cloud-token: ${{ secrets.ATLAS_CLOUD_TOKEN }}
        version: 'beta'
    - uses: ariga/atlas-action/schema/plan@v1
      with:
        env: demo
        dev-url: '${{ secrets.DATABRICKS_DEV_URL }}'

This first plan section occurs when the pull request is opened.

Here, we ensure that Atlas is installed and use your Atlas Cloud token for authentication.

Then, Atlas creates the migration plan that takes our target database to the desired schema state defined in our "demo" environment in the atlas.hcl file. This plan is then linted and both the plan and lint report are commented on the PR, along with a link to view the plan in your Atlas Cloud repository.

push:
  if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }}
  runs-on: ubuntu-latest
  env:
    GITHUB_TOKEN: ${{ github.token }}
  steps:
    - uses: actions/checkout@v3
      with:
        fetch-depth: 0
    - uses: ariga/setup-atlas@v0
      with:
        cloud-token: ${{ secrets.ATLAS_CLOUD_TOKEN }}
        version: 'beta'
    - uses: ariga/atlas-action/schema/plan/approve@v1
      id: plan-approve
      with:
        env: demo
        dev-url: '${{ secrets.DATABRICKS_DEV_URL }}'
    - uses: ariga/atlas-action/schema/push@v1
      with:
        env: demo
        dev-url: '${{ secrets.DATABRICKS_DEV_URL }}'
    - uses: ariga/atlas-action/schema/apply@v1
      with:
        env: demo
        dev-url: '${{ secrets.DATABRICKS_DEV_URL }}'
        tx-mode: none

This next push section occurs when the pull request is merged to the main branch.

We start by re-ensuring that Atlas is properly installed and Atlas Cloud is authenticated.

Next, the migration plan is marked "approved" and the new schema pushed to the Atlas Cloud repository.

Finally, the migration plan is applied to the target database.

Versioned workflow in CI/CD

First, set up an Atlas Cloud repository by running:

atlas migrate push databricks-demo --dev-url "$DATABRICKS_DEV_URL"

Replace databricks-demo with your Atlas Cloud repository name.

When you create a new migration file locally, push the updated desired schema and migrations directory to a new branch in your GitHub repository and make a pull request. This will trigger the workflow file.

Workflow file

Copy the .github/workflows/ci-atlas-versioned.yaml file into your repository, replacing databricks-demo with your Atlas Cloud repository name.

Let's go over the jobs.

- uses: ariga/setup-atlas@v0
  with:
    cloud-token: '${{ secrets.ATLAS_CLOUD_TOKEN }}'
    version: 'beta'

Here, we ensure that Atlas is installed and use the Atlas Cloud token for authentication.

- uses: ariga/atlas-action/migrate/lint@v1
  with:
    dir: 'file://migrations'
    dev-url: '${{ secrets.DATABRICKS_DEV_URL }}'
    dir-name: 'databricks-demo'
  env:
    GITHUB_TOKEN: '${{ github.token }}'

Next, we lint the new migration file using Atlas's analyzers and checks to ensure that the migration is safe to apply. The results are reported on the PR when complete and resported to Atlas Cloud.

Note

You can define custom rules based on your organization's policy that the linter will also verify. Read more in our documentation.

- uses: ariga/atlas-action/migrate/push@v1
  if: github.ref == 'refs/heads/main'
  with:
    dir: 'file://migrations'
    dev-url:  '${{ secrets.DATABRICKS_DEV_URL }}'
    dir-name: 'databricks-demo'
- uses: ariga/atlas-action/migrate/apply@v1
  if: github.ref == 'refs/heads/main'
  with:
    dir: 'file://migrations'
    url: '${{ secrets.DATABRICKS_URL }}'
    tx-mode: none

When the pull request is merged to the main branch, two actions occur:

1. Atlas pushes the updated migration directory to the Atlas Cloud repository
2. Atlas applies the new migration to the target database

Tip

It is also possible to automate the migrate diff command at the start, so all you have to do is push the new desired schema, similar to the declarative workflow.

Go to our GitHub Actions documentation to learn more.

Note that this may limit migrate diff variations, such as adding keywords to the migration file names.

Next steps

If you have any further questions, feel free to ask via the Intercom button at the bottom corner of our website, the Ask AI at the top of our docs page, or our Discord server.

References

• Databricks with Atlas guide
• HCL reference for Databricks schema definitions
• Atlas CI/CD on GitHub
• Declarative vs Versioned Workflows article
• Declarative workflow docs
• Versioned workflow docs