Database Migrations & Schema Evolution¶

# alembic/env.py
from logging.config import fileConfig
from sqlalchemy import engine_from_config
from sqlalchemy import pool
from alembic import context
import os
import sys

# Add your model's MetaData object here
from myapp.models import Base
target_metadata = Base.metadata

# Alembic Config object
config = context.config

# Interpret the config file for Python logging
if config.config_file_name is not None:
    fileConfig(config.config_file_name)

def run_migrations_offline():
    """Run migrations in 'offline' mode."""
    url = config.get_main_option("sqlalchemy.url")
    context.configure(
        url=url,
        target_metadata=target_metadata,
        literal_binds=True,
        dialect_opts={"paramstyle": "named"},
    )

    with context.begin_transaction():
        context.run_migrations()

def run_migrations_online():
    """Run migrations in 'online' mode."""
    connectable = engine_from_config(
        config.get_section(config.config_ini_section),
        prefix="sqlalchemy.",
        poolclass=pool.NullPool,
    )

    with connectable.connect() as connection:
        context.configure(
            connection=connection, target_metadata=target_metadata
        )

        with context.begin_transaction():
            context.run_migrations()

if context.is_offline_mode():
    run_migrations_offline()
else:
    run_migrations_online()

# Migration: Add new column without breaking existing code
"""Add user preferences table

Revision ID: 001_add_user_preferences
Revises: 
Create Date: 2024-01-15 10:00:00.000000

"""
from alembic import op
import sqlalchemy as sa

revision = '001_add_user_preferences'
down_revision = None
branch_labels = None
depends_on = None

def upgrade():
    # Step 1: Add new table
    op.create_table('user_preferences',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('user_id', sa.Integer(), nullable=False),
        sa.Column('theme', sa.String(50), nullable=True, default='light'),
        sa.Column('notifications', sa.Boolean(), nullable=True, default=True),
        sa.Column('created_at', sa.DateTime(timezone=True), nullable=False),
        sa.Column('updated_at', sa.DateTime(timezone=True), nullable=False),
        sa.ForeignKeyConstraint(['user_id'], ['users.id'], ondelete='CASCADE'),
        sa.PrimaryKeyConstraint('id')
    )

    # Step 2: Create indexes
    op.create_index('idx_user_preferences_user_id', 'user_preferences', ['user_id'])

    # Step 3: Add default preferences for existing users
    op.execute("""
        INSERT INTO user_preferences (user_id, theme, notifications, created_at, updated_at)
        SELECT id, 'light', true, NOW(), NOW()
        FROM users
        WHERE NOT EXISTS (
            SELECT 1 FROM user_preferences WHERE user_id = users.id
        )
    """)

def downgrade():
    # Step 1: Drop indexes
    op.drop_index('idx_user_preferences_user_id', table_name='user_preferences')

    # Step 2: Drop table
    op.drop_table('user_preferences')

# Migration: Rename column with zero downtime
"""Rename username to user_handle

Revision ID: 002_rename_username_to_user_handle
Revises: 001_add_user_preferences
Create Date: 2024-01-15 11:00:00.000000

"""
from alembic import op
import sqlalchemy as sa

revision = '002_rename_username_to_user_handle'
down_revision = '001_add_user_preferences'
branch_labels = None
depends_on = None

def upgrade():
    # Step 1: Add new column
    op.add_column('users', sa.Column('user_handle', sa.String(50), nullable=True))

    # Step 2: Copy data from old column to new column
    op.execute("UPDATE users SET user_handle = username WHERE user_handle IS NULL")

    # Step 3: Make new column NOT NULL (after data is copied)
    op.alter_column('users', 'user_handle', nullable=False)

    # Step 4: Create unique index on new column
    op.create_index('idx_users_user_handle', 'users', ['user_handle'], unique=True)

    # Note: Don't drop old column yet - wait for application code to be updated

def downgrade():
    # Step 1: Drop new column index
    op.drop_index('idx_users_user_handle', table_name='users')

    # Step 2: Drop new column
    op.drop_column('users', 'user_handle')

# tests/migrations/test_migration_validation.py
import pytest
import psycopg2
from sqlalchemy import create_engine, text
from alembic import command
from alembic.config import Config

class TestMigrationValidation:
    """Test that migrations work correctly and don't break data."""

    @pytest.fixture
    def db_engine(self):
        """Create test database engine."""
        return create_engine("postgresql://testuser:testpass@localhost:5432/testdb")

    @pytest.fixture
    def alembic_cfg(self):
        """Create Alembic configuration."""
        config = Config("alembic.ini")
        config.set_main_option("sqlalchemy.url", "postgresql://testuser:testpass@localhost:5432/testdb")
        return config

    def test_migration_upgrade(self, db_engine, alembic_cfg):
        """Test that migration upgrade works."""
        # Run migration
        command.upgrade(alembic_cfg, "head")

        # Validate schema
        with db_engine.connect() as conn:
            # Check tables exist
            result = conn.execute(text("""
                SELECT table_name 
                FROM information_schema.tables 
                WHERE table_schema = 'public'
                ORDER BY table_name
            """))
            tables = [row[0] for row in result]
            assert 'users' in tables
            assert 'user_preferences' in tables
            assert 'telemetry' in tables

    def test_migration_downgrade(self, db_engine, alembic_cfg):
        """Test that migration downgrade works."""
        # Upgrade first
        command.upgrade(alembic_cfg, "head")

        # Then downgrade
        command.downgrade(alembic_cfg, "base")

        # Validate schema is clean
        with db_engine.connect() as conn:
            result = conn.execute(text("""
                SELECT table_name 
                FROM information_schema.tables 
                WHERE table_schema = 'public'
            """))
            tables = [row[0] for row in result]
            assert len(tables) == 0

    def test_data_integrity(self, db_engine, alembic_cfg):
        """Test that data integrity is maintained during migration."""
        # Insert test data before migration
        with db_engine.connect() as conn:
            conn.execute(text("""
                INSERT INTO users (username, email) 
                VALUES ('testuser', 'test@example.com')
            """))
            conn.commit()

        # Run migration
        command.upgrade(alembic_cfg, "head")

        # Validate data is still there
        with db_engine.connect() as conn:
            result = conn.execute(text("SELECT COUNT(*) FROM users"))
            count = result.scalar()
            assert count == 1

    def test_geospatial_data(self, db_engine, alembic_cfg):
        """Test that geospatial data is handled correctly."""
        # Run migration
        command.upgrade(alembic_cfg, "head")

        # Insert geospatial test data
        with db_engine.connect() as conn:
            conn.execute(text("""
                INSERT INTO telemetry (user_id, sensor_id, timestamp, location)
                VALUES (1, 'sensor_001', NOW(), ST_GeomFromText('POINT(-122.4194 37.7749)', 4326))
            """))
            conn.commit()

        # Validate geospatial data
        with db_engine.connect() as conn:
            result = conn.execute(text("""
                SELECT ST_AsText(location), ST_SRID(location)
                FROM telemetry
                WHERE sensor_id = 'sensor_001'
            """))
            row = result.fetchone()
            assert row[0] == 'POINT(-122.4194 37.7749)'
            assert row[1] == 4326

# scripts/validate_migration.py
import os
import sys
import psycopg2
from sqlalchemy import create_engine, text

def validate_migration():
    """Validate that migration completed successfully."""
    db_url = os.getenv("DATABASE_URL")
    if not db_url:
        print("ERROR: DATABASE_URL not set")
        sys.exit(1)

    engine = create_engine(db_url)

    with engine.connect() as conn:
        # Check that all tables exist
        result = conn.execute(text("""
            SELECT table_name 
            FROM information_schema.tables 
            WHERE table_schema = 'public'
            ORDER BY table_name
        """))
        tables = [row[0] for row in result]
        expected_tables = ['users', 'user_preferences', 'telemetry']

        for table in expected_tables:
            if table not in tables:
                print(f"ERROR: Table {table} not found")
                sys.exit(1)

        # Check row counts
        result = conn.execute(text("SELECT COUNT(*) FROM users"))
        user_count = result.scalar()
        print(f"Users: {user_count}")

        # Check geospatial data integrity
        result = conn.execute(text("""
            SELECT COUNT(*) 
            FROM telemetry 
            WHERE location IS NOT NULL
        """))
        geo_count = result.scalar()
        print(f"Geospatial records: {geo_count}")

        # Check SRID consistency
        result = conn.execute(text("""
            SELECT DISTINCT ST_SRID(location) as srid
            FROM telemetry 
            WHERE location IS NOT NULL
        """))
        srids = [row[0] for row in result]
        if len(srids) > 1:
            print(f"WARNING: Multiple SRIDs found: {srids}")
        else:
            print(f"SRID consistency: {srids[0] if srids else 'No geospatial data'}")

        print("Migration validation completed successfully")

if __name__ == "__main__":
    validate_migration()