database migrations setup

Understanding Database Migrations

Database migrations are controlled sets of changes applied to a database schema over time. They allow developers to evolve the database structure (adding tables, modifying columns, setting up indexes, etc.) without disrupting existing data or applications that rely on the database.

For example, consider a PostgreSQL database for a blogging platform with an initial posts table:

CREATE TABLE posts (
    id SERIAL PRIMARY KEY,
    title VARCHAR(255) NOT NULL,
    content TEXT NOT NULL,
    published_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);

Suppose you need to add a slug column to improve SEO-friendly URLs. A database migration would involve:

Writing a script to alter the table:

ALTER TABLE posts ADD COLUMN slug VARCHAR(255) UNIQUE;

Applying the script to the database in a controlled manner, ensuring all environments (development, staging, production) stay in sync.

Why Use Migrations?

Version Control: Keep a history of schema changes.
Collaboration: Multiple developers can work on the database schema concurrently.
Rollback Capability: Revert changes if something goes wrong.
Consistency: Ensure all environments reflect the same schema state.

Introducing Alembic

Alembic is a lightweight database migration tool for Python, designed to work with SQLAlchemy. It manages database schema changes in a systematic and version-controlled way.

Key Features:

Versioning: Keeps track of migrations via version numbers.
Autogeneration: Automatically generates migration scripts by comparing your models to the current database schema.
Upgrade/Downgrade: Applies or reverts migrations as needed.
CLI Interface: Provides command-line tools for managing migrations.

Using Alembic with SQLModel for Migrations

SQLModel is a library that combines SQLAlchemy and Pydantic, making it easier to work with SQL databases in Python.

Here's how to use Alembic with SQLModel for database migrations:

1. Define SQLModel Models

Create models.py:

from sqlmodel import SQLModel, Field
from typing import Optional

class Post(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    title: str
    content: str
    published_at: Optional[str]

2. Initialize Alembic

Run:

alembic init alembic

This command creates an alembic directory with necessary configurations.

3. Configure Alembic

Edit alembic.ini to set your database URL:

sqlalchemy.url = postgresql+psycopg2://username:password@localhost/mydatabase

Modify alembic/env.py to import your models and set up target_metadata:

from models import SQLModel

target_metadata = SQLModel.metadata

4. Create the Initial Migration

Generate a new migration script:

alembic revision --autogenerate -m "Initial migration"

This command compares your models with the current database schema and generates a migration script in alembic/versions/.

Review the Migration Script:

def upgrade():
    op.create_table(
        'post',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('title', sa.String(), nullable=False),
        sa.Column('content', sa.String(), nullable=False),
        sa.Column('published_at', sa.String(), nullable=True),
        sa.PrimaryKeyConstraint('id')
    )

7. Apply the Migration

Run:

alembic upgrade head

This command applies all migrations up to the latest (head) version.

5. Update Models and Create New Migrations

Suppose you want to add a slug column to the Post model.

Update models.py:

class Post(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    title: str
    slug: str  # New field
    content: str
    published_at: Optional[str]

Generate a New Migration:

alembic revision --autogenerate -m "Add slug to Post"

Review the New Migration Script:

def upgrade():
    op.add_column('post', sa.Column('slug', sa.String(), nullable=False))
    op.create_unique_constraint(None, 'post', ['slug'])

Apply the Migration:

alembic upgrade head

6. Rolling Back Migrations

If you need to undo the last migration:

alembic downgrade -1

This command reverts the last applied migration.

7. Handling Complex Migrations

For more complex changes, like renaming a column or modifying data during migration, you may need to edit the migration script manually.

Example: Renaming a Column

Suppose you want to rename published_at to published_on.

Update models.py:

class Post(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    title: str
    slug: str
    content: str
    published_on: Optional[str]  # Renamed field

Generate the Migration:

alembic revision --autogenerate -m "Rename published_at to published_on"

Edit the Migration Script Manually:

Alembic's autogeneration may not detect column renames. Manually adjust the script:

def upgrade():
    op.alter_column('post', 'published_at', new_column_name='published_on')

Complete Example: Putting It All Together

models.py

from sqlmodel import SQLModel, Field
from typing import Optional

class User(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    username: str
    email: str

class Post(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    title: str
    slug: str
    content: str
    author_id: int = Field(foreign_key="user.id")
    published_on: Optional[str]

Generate Migrations:

Create Users Table

alembic revision --autogenerate -m "Create users table"
alembic upgrade head

Add Foreign Key to Posts

After updating the Post model with author_id, generate and apply the migration:
```
alembic revision --autogenerate -m "Add author_id to posts"
alembic upgrade head
```

main.py

from sqlmodel import Session, create_engine
from models import SQLModel, User, Post

engine = create_engine("postgresql+psycopg2://username:password@localhost/mydatabase")

def create_user_and_post():
    with Session(engine) as session:
        user = User(username="johndoe", email="john@example.com")
        session.add(user)
        session.commit()
        session.refresh(user)

        post = Post(
            title="My First Post",
            slug="my-first-post",
            content="Hello, World!",
            author_id=user.id
        )
        session.add(post)
        session.commit()

if __name__ == "__main__":
    create_user_and_post()

Summary

Database Migrations: Essential for evolving your database schema in a controlled, versioned manner.
Alembic: A powerful tool that integrates with SQLAlchemy (and by extension, SQLModel) to manage migrations.
SQLModel: Simplifies working with SQL databases by combining SQLAlchemy and Pydantic features.
Workflow:
1. Define or update models in models.py.
2. Generate migrations using Alembic.
3. Review and apply migrations.
4. Use the updated models in your application code.

By integrating Alembic with SQLModel, you can efficiently manage your PostgreSQL database schema, ensuring consistency across development, testing, and production environments.

Additional Tips

Version Control: Always commit your migration scripts to version control (e.g., Git).
Testing Migrations: Before applying to production, test migrations in a staging environment.
Handling Data Migrations: For data transformations, include data manipulation code within your migration scripts.
Staying Updated: Keep an eye on updates to Alembic and SQLModel for new features and best practices.

capstone-UniGate / unigate