Skip to main content

Automatic Migration Planning For SQLAlchemy

TL;DR​

  • SQLAlchemy is a Python SQL toolkit and Object Relational Mapper.
  • Atlas is an open-source tool for inspecting, planning, linting and executing schema changes to your database.
  • Developers using SQLAlchemy can use Atlas to automatically plan schema migrations for them, based on the desired state of their schema instead of crafting them by hand.

Prerequisites​

If you don't have a SQLAlchemy project handy, check out the SQLAlchemy quick start page.

Using the Atlas SQLAlchemy Provider​

In this guide, we will use the SQLAlchemy Atlas Provider to automatically plan schema migrations for a SQLAlchemy project.

To use Atlas with SQLAlchemy, there are two modes in which the Atlas SQLAlchemy Provider can load your schema.

  • Standalone Mode - If all of your SQLAlchemy models exist in a single module, you can use the provider directly to load your SQLAlchemy schema into Atlas.
  • Script Mode - Use Atlas as a Python script to load and manage your SQLAlchemy schema in Atlas.

In this getting-started guide, we will be using Standalone Mode, but both modes require the same installation.

Installation​

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

curl -sSf https://atlasgo.sh | sh

Install the provider by running:

pip install atlas-provider-sqlalchemy
TL;DR

The SQLAlchemy Atlas Provider works by importing your SQLAlchemy models and extracting the schema from them.

Therefore, you will need to run the provider from within your project's Python environment.

Example Usage​

We will demonstrate the provider with the versioned migrations workflow, where each change to the database is versioned and recorded in a migration file. You can use the atlas migrate diff command to automatically generate a migration file that will migrate the database from its latest revision to the current SQLAlchemy schema.

In our example, we have a directory named models that contains all of our SQLAlchemy models:

models/models.py
from sqlalchemy import ForeignKey, String
from sqlalchemy.orm import Mapped, mapped_column, relationship, DeclarativeBase


class Base(DeclarativeBase):
    pass


class User(Base):
    __tablename__ = "user_account"
    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str] = mapped_column(String(30))
    fullname: Mapped[str | None] = mapped_column(String(30))
    addresses: Mapped[list["Address"]] = relationship(
        back_populates="user", cascade="all, delete-orphan"
    )


class Address(Base):
    __tablename__ = "address"
    id: Mapped[int] = mapped_column(primary_key=True)
    email_address: Mapped[str] = mapped_column(String(30))
    user_id: Mapped[int] = mapped_column(ForeignKey("user_account.id"))
    user: Mapped["User"] = relationship(back_populates="addresses")

To configure the provider, create an atlas.hcl file with the following configuration:

atlas.hcl
data "external_schema" "sqlalchemy" {
    program = [
        "atlas-provider-sqlalchemy",
        "--path", "./models",
        "--dialect", "postgresql"
    ]
}

env "sqlalchemy" {
    src = data.external_schema.sqlalchemy.url
    dev = "docker://postgres/16/dev?search_path=public"
    migration {
        dir = "file://migrations"
    }
    format {
        migrate {
            diff = "{{ sql . "  " }}"
        }
    }
}

New, generate a new migration using the following command:

atlas migrate diff --env sqlalchemy

Running this command will generate files similar to this in the migrations directory:

migrations
|-- 20230918143104.sql
`-- atlas.sum

0 directories, 2 files

Examining the contents of 20230918143104.sql:

migrations/20230918143104.sql
-- Create "user_account" table
CREATE TABLE `user_account` (
  `id` int NOT NULL AUTO_INCREMENT,
  `name` varchar(30) NOT NULL,
  `fullname` varchar(30) NULL,
  PRIMARY KEY (`id`)
) CHARSET utf8mb4 COLLATE utf8mb4_0900_ai_ci;
-- Create "address" table
CREATE TABLE `address` (
  `id` int NOT NULL AUTO_INCREMENT,
  `email_address` varchar(30) NOT NULL,
  `user_id` int NOT NULL,
  PRIMARY KEY (`id`),
  INDEX `user_id` (`user_id`),
  CONSTRAINT `address_ibfk_1` FOREIGN KEY (`user_id`) REFERENCES `user_account` (`id`) ON UPDATE NO ACTION ON DELETE NO ACTION
) CHARSET utf8mb4 COLLATE utf8mb4_0900_ai_ci;

Amazing! Atlas automatically generated a migration file that will create the user_account and address tables in our database.

Next, alter the user_account entity to add a new age property:

    fullname: Mapped[Optional[str]] = mapped_column(String(30))
+   age: Mapped[Optional[int]] = mapped_column()
    addresses: Mapped[List["Address"]] = relationship(
        back_populates="user", cascade="all, delete-orphan"
    )

Re-run this command:

atlas migrate diff --env sqlalchemy

Observe a new migration file is generated in the migration directory:

-- Modify "user_account" table
ALTER TABLE `user_account` ADD COLUMN `age` int NULL;

Conclusion​

In this guide we demonstrated how projects using SQLAlchemy can use Atlas to automatically plan schema migrations based only on their data model. To learn more about executing migrations against your production database, read the documentation for the migrate apply command.

Have questions? Feedback? Find our team on our Discord server