Migrate CockroachDB Schemas with Alembic

On this page

This page guides you through a series of simple database schema changes using the Alembic schema migration module with a simple Python application built on SQLAlchemy and CockroachDB.

For a detailed tutorial about using Alembic, see the Alembic documentation site.

For information about specific migration tasks, see Alembic's Cookbook.

Before you begin

Before you begin the tutorial, install CockroachDB.

Step 1. Start a cluster and create a database

Start a demo cluster:
```
$ cockroach demo --no-example-database
```
This command creates a virtual cluster and opens a SQL shell to that cluster.

Note:

Leave this terminal window open for the duration of the tutorial. Closing the window will destroy the cluster and erase all data in it.
Create the bank database:
```
> CREATE DATABASE bank;
```

Step 2. Get the application code

Open a new terminal, and clone the example-app-python-sqlalchemy GitHub repository:
```
$ git clone git@github.com:cockroachlabs/example-app-python-sqlalchemy.git
```

Step 3. Install and initialize Alembic

Navigate to the example-app-python-sqlalchemy project directory, and run the following commands to create and start a virtual environment:
```
$ python3 -m venv env
```
```
$ source env/bin/activate
```
Install the alembic, sqlalchemy-cockroachdb, and psycopg2 modules to the virtual environment:
```
$ pip install sqlalchemy-cockroachdb psycopg2-binary alembic
```
The sqlalchemy-cockroachdb and psycopg2-binary modules are required to use the CockroachDB adapter that the app uses to run transactions against a CockroachDB cluster.

alembic includes the sqlalchemy module, which is a primary dependency of the example-app-python-sqlalchemy sample app. The alembic install also includes the alembic command line tool, which we use throughout the tutorial to manage migrations.

Use the alembic command-line tool to initialize Alembic for the project:

$ alembic init alembic

Creating directory /path/example-app-python-
sqlalchemy/alembic ...  done
Creating directory /path/example-app-python-
sqlalchemy/alembic/versions ...  done
Generating /path/example-app-python-
sqlalchemy/alembic/script.py.mako ...  done
Generating /path/example-app-python-
sqlalchemy/alembic/env.py ...  done
Generating /path/example-app-python-
sqlalchemy/alembic/README ...  done
Generating /path/example-app-python-
sqlalchemy/alembic.ini ...  done
Please edit configuration/connection/logging settings in
'/path/example-app-python-sqlalchemy/alembic.ini' before
proceeding.

This command creates a migrations directory called alembic. This directory will contain the files that specify the schema migrations for the app.

The command also creates a properties file called alembic.ini at the top of the project directory.

Open alembic.ini and update the sqlalchemy.url property to specify the correct connection string to your database:

For example:
```
sqlalchemy.url = cockroachdb://demo:demo72529@127.0.0.1:26257/bank?sslmode=require
```
Note:

You must use the cockroachdb:// prefix in the connection string for SQLAlchemy to make sure the CockroachDB dialect is used. Using the postgresql:// URL prefix to connect to your CockroachDB cluster will not work.

Step 4. Create and run a migration script

Use the alembic command-line tool to create the first migration script:

$ alembic revision -m "create accounts table"

Generating /path/example-app-python-sqlalchemy/alembic/versions/ad72c7ec8b22_create_accounts_table.py ...  done

Open the newly-created migration file (alembic/versions/ad72c7ec8b22_create_accounts_table.py, in this case), and edit the upgrade() and downgrade() functions to read as follows:
```
def upgrade():
    op.create_table(
        'accounts',
        sa.Column('id', sa.dialects.postgresql.UUID, primary_key=True),
        sa.Column('balance', sa.Integer),
    )

def downgrade():
    op.drop_table('accounts')
```
Running this migration creates the accounts table, with an id column and a balance column.

Note that this file also specifies an operation for "downgrading" the migration. In this case, downgrading will drop the accounts table, effectively reversing the schema changes of the migration.
Use the alembic tool to run this first migration:
```
$ alembic upgrade head
```
```
INFO  [alembic.runtime.migration] Context impl CockroachDBImpl.
INFO  [alembic.runtime.migration] Will assume non-transactional DDL.
INFO  [alembic.runtime.migration] Running upgrade  -> ad72c7ec8b22, create accounts table
```
Specifying head runs the latest migration. This migration will create the accounts table. It will also create a table called alembic_version, which tracks the current migration version of the database.

Step 5. Verify the migration

Open the terminal with the SQL shell to your demo cluster, and verify that the table was successfully created:

> USE bank;

> SHOW TABLES;

  schema_name |   table_name    | type  | owner | estimated_row_count | locality
--------------+-----------------+-------+-------+---------------------+-----------
  public      | accounts        | table | demo  |                   0 | NULL
  public      | alembic_version | table | demo  |                   1 | NULL
(2 rows)

> SELECT * FROM alembic_version;

  version_num
----------------
  ad72c7ec8b22
(1 row)

> SHOW COLUMNS FROM accounts;

  column_name | data_type | is_nullable | column_default | generation_expression |  indices  | is_hidden
--------------+-----------+-------------+----------------+-----------------------+-----------+------------
  id          | UUID      |    false    | NULL           |                       | {primary} |   false
  balance     | INT8      |    true     | NULL           |                       | {primary} |   false
(2 rows)

In a different terminal, set the DATABASE_URL environment variable to the connection string for your cluster:
```
$ export DATABASE_URL=cockroachdb://demo:demo72529@127.0.0.1:26257/bank?sslmode=require
```
The sample app reads in DATABASE_URL as the connection string to the database.

Run the app to insert, update, and delete rows of data:

$ python main.py

Creating new accounts...
Created new account with id e9b4a9da-fbbb-40de-8c44-60c5c741764d and balance 93911.
Created new account with id 34a6b5d6-0f08-4435-89cb-c7fa30037926 and balance 989744.
...
Created new account with id 18a7a209-72c3-48b6-986c-2631fff38274 and balance 969474.
Created new account with id 68e73209-fe2e-42db-a54e-c9d101990cdc and balance 382471.
Random account balances:
Account 9acbf774-3e22-4d75-aee0-37e63d3b1ab6: 403963
Account 82451815-3a87-4d67-a9b0-7766726abd31: 315597
Transferring 201981 from account 9acbf774-3e22-4d75-aee0-37e63d3b1ab6 to account 82451815-3a87-4d67-a9b0-7766726abd31...
Transfer complete.
New balances:
Account 9acbf774-3e22-4d75-aee0-37e63d3b1ab6: 201982
Account 82451815-3a87-4d67-a9b0-7766726abd31: 517578
Deleting existing accounts...
Deleted account 13d1b940-9a7b-47d6-b719-6a2b49a3b08c.
Deleted account 6958f8f9-4d38-424c-bf41-5673f20169b1.
Deleted account c628bd7f-3054-4cd6-b2c9-8c2e3def1720.
Deleted account f4268300-6d0a-4d6e-9489-ad30f215d1ad.
Deleted account feae4e4a-c003-4c29-b672-5422438a885b.

Step 6. Add additional migrations

Suppose you want to add a new computed column to the accounts table that tracks which accounts are overdrawn.

Create a new migration with the alembic tool:

$ alembic revision -m "add overdrawn column"

Generating /path/example-app-python-sqlalchemy/alembic/versions/fd88c68af7b5_add_overdrawn_column.py ...  done

Open the migration file (alembic/versions/fd88c68af7b5_add_overdrawn_column.py), update the imports, and edit the upgrade() and downgrade() functions:

from alembic import op
from sqlalchemy import Column, Boolean, Computed

...

def upgrade():
    op.add_column('accounts', sa.Column('overdrawn', Boolean, Computed('CASE WHEN balance < 0 THEN True ELSE False END')))

def downgrade():
    op.drop_column('accounts', 'overdrawn')

Use the alembic tool to run the migration.

Because this is the latest migration, you can specify head, or you can use the migration's ID (fd88c68af7b5):

$ alembic upgrade fd88c68af7b5

INFO  [alembic.runtime.migration] Context impl CockroachDBImpl.
INFO  [alembic.runtime.migration] Will assume non-transactional DDL.
INFO  [alembic.runtime.migration] Running upgrade ad72c7ec8b22 -> fd88c68af7b5, add_overdrawn_column

In the terminal with the SQL shell to your demo cluster, verify that the column was successfully created:

> SHOW COLUMNS FROM accounts;

  column_name | data_type | is_nullable | column_default |             generation_expression              |  indices  | is_hidden
--------------+-----------+-------------+----------------+------------------------------------------------+-----------+------------
  id          | UUID      |    false    | NULL           |                                                | {primary} |   false
  balance     | INT8      |    true     | NULL           |                                                | {primary} |   false
  overdrawn   | BOOL      |    true     | NULL           | CASE WHEN balance < 0 THEN true ELSE false END | {primary} |   false
(3 rows)

> SELECT * FROM accounts;

                   id                  | balance | overdrawn
---------------------------------------+---------+------------
  01894212-7f32-4e4c-b855-146630d928bc |  548554 |   false
  033131cf-7c42-4021-9a53-f8a7597ec853 |  828874 |   false
  041a2c5d-0bce-4ed4-a91d-a9e3a6e06632 |  768526 |   false
  080be3a3-40f8-40c6-a0cc-a61c108db3f5 |  599729 |   false
  08503245-ba1a-4255-8ca7-22b3688e69dd |    7962 |   false
  ...

The changes will also be reflected in the alembic_version table.

> SELECT * FROM alembic_version;

  version_num
----------------
  fd88c68af7b5
(1 row)

Execute Raw SQL with Alembic

While Alembic supports most SQL operations, you can always execute raw SQL using the execute() operation.

Tip:

Executing DDL statements as raw SQL can be particularly helpful when using SQL syntax for DDL statements specific to CockroachDB, like ALTER TABLE ... ALTER PRIMARY KEY or ALTER TABLE ... SET LOCALITY statements.

For example, the raw SQL for the second migration would look something like this:

ALTER TABLE accounts ADD COLUMN overdrawn BOOLEAN AS (
  CASE
    WHEN balance < 0 THEN True
    ELSE False
  END
) STORED;

To make the second migration use raw SQL instead of Alembic operations, open alembic/versions/fd88c68af7b5_add_overdrawn_column.py, and edit the upgrade() function to use execute() instead of the operation-specific function:

def upgrade():
    op.execute(text("""ALTER TABLE accounts ADD COLUMN overdrawn BOOLEAN AS (
                        CASE
                            WHEN balance < 0 THEN True
                            ELSE False
                        END
                        ) STORED;"""))

Before running this migration, downgrade the original migration:

$ alembic downgrade -1

INFO  [alembic.runtime.migration] Context impl CockroachDBImpl.
INFO  [alembic.runtime.migration] Will assume non-transactional DDL.
INFO  [alembic.runtime.migration] Running downgrade fd88c68af7b5 -> ad72c7ec8b22, add_overdrawn_column

Then, in the SQL shell to the demo cluster, verify that the overdrawn column has been dropped from the table:

> SHOW COLUMNS FROM accounts;

  column_name | data_type | is_nullable | column_default | generation_expression |  indices  | is_hidden
--------------+-----------+-------------+----------------+-----------------------+-----------+------------
  id          | UUID      |    false    | NULL           |                       | {primary} |   false
  balance     | INT8      |    true     | NULL           |                       | {primary} |   false
(2 rows)

Now, run the updated migration script:

$ alembic upgrade fd88c68af7b5

INFO  [alembic.runtime.migration] Context impl CockroachDBImpl.
INFO  [alembic.runtime.migration] Will assume non-transactional DDL.
INFO  [alembic.runtime.migration] Running upgrade ad72c7ec8b22 -> fd88c68af7b5, add_overdrawn_column

And verify that the column has been added to the table:

> SHOW COLUMNS FROM accounts;

  column_name | data_type | is_nullable | column_default |             generation_expression              |  indices  | is_hidden
--------------+-----------+-------------+----------------+------------------------------------------------+-----------+------------
  id          | UUID      |    false    | NULL           |                                                | {primary} |   false
  balance     | INT8      |    true     | NULL           |                                                | {primary} |   false
  overdrawn   | BOOL      |    true     | NULL           | CASE WHEN balance < 0 THEN true ELSE false END | {primary} |   false
(3 rows)

Auto-generate a Migration

Alembic can automatically generate migrations, based on changes to the models in your application source code.

Let's use the same example overdrawn computed column from above.

First, downgrade the fd88c68af7b5 migration:

$ alembic downgrade -1

INFO  [alembic.runtime.migration] Context impl CockroachDBImpl.
INFO  [alembic.runtime.migration] Will assume non-transactional DDL.
INFO  [alembic.runtime.migration] Running downgrade fd88c68af7b5 -> ad72c7ec8b22, add_overdrawn_column

> SHOW COLUMNS FROM accounts;

  column_name | data_type | is_nullable | column_default | generation_expression |  indices  | is_hidden
--------------+-----------+-------------+----------------+-----------------------+-----------+------------
  id          | UUID      |    false    | NULL           |                       | {primary} |   false
  balance     | INT8      |    true     | NULL           |                       | {primary} |   false
(2 rows)

Delete the old migration file:

rm alembic/versions/fd88c68af7b5_add_overdrawn_column.py

Open the models.py file in the app's project, and add the overdrawn column to the Account class definition:

from sqlalchemy import Column, Integer, Boolean, Computed

...

class Account(Base):
    """The Account class corresponds to the "accounts" database table.
    """
    __tablename__ = 'accounts'
    id = Column(UUID(as_uuid=True), primary_key=True)
    balance = Column(Integer)
    overdrawn = Column('overdrawn', Boolean, Computed('CASE WHEN balance < 0 THEN True ELSE False END'))

Then, open the alembic/env.py file, and add the following import to the top of the file:

from ..models import Base

And update the variable target_metadata to read as follows:

target_metadata = Base.metadata

These two lines import the database model metadata from the app.

Use the alembic command-line tool to auto-generate the migration from the models defined in the app:

$ alembic revision --autogenerate -m "add overdrawn column"

INFO  [alembic.runtime.migration] Context impl CockroachDBImpl.
INFO  [alembic.runtime.migration] Will assume non-transactional DDL.
INFO  [alembic.autogenerate.compare] Detected added column 'accounts.overdrawn'
  Generating /path/example-app-python-sqlalchemy/alembic/versions/44fa7043e441_add_overdrawn_column.py ...  done

Alembic creates a new migration file (44fa7043e441_add_overdrawn_column.py, in this case).

If you open this file, you'll see that it looks very similar to the one you manually created earlier in the tutorial.

...
def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('accounts', sa.Column('overdrawn', sa.Boolean(), sa.Computed('CASE WHEN balance < 0 THEN True ELSE False END', ), nullable=True))
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('accounts', 'overdrawn')
    # ### end Alembic commands ###

Run the migration:

$ alembic upgrade 44fa7043e441

INFO  [alembic.runtime.migration] Context impl CockroachDBImpl.
INFO  [alembic.runtime.migration] Will assume non-transactional DDL.
INFO  [alembic.runtime.migration] Running upgrade ad72c7ec8b22 -> 44fa7043e441, add overdrawn column

Verify that the new column exists in the accounts table:

> SHOW COLUMNS FROM accounts;

  column_name | data_type | is_nullable | column_default |             generation_expression              |  indices  | is_hidden
--------------+-----------+-------------+----------------+------------------------------------------------+-----------+------------
  id          | UUID      |    false    | NULL           |                                                | {primary} |   false
  balance     | INT8      |    true     | NULL           |                                                | {primary} |   false
  overdrawn   | BOOL      |    true     | NULL           | CASE WHEN balance < 0 THEN true ELSE false END | {primary} |   false
(3 rows)

Report Issues with Alembic and CockroachDB

If you run into problems, please file an issue in the alembic repository, including the following details about the environment where you encountered the issue:

CockroachDB version (cockroach version)
Alembic version
Operating system
Steps to reproduce the behavior

Pricing

Contact us

Sign In