Migration Examples - Manticore CockroachDB Client
Description: A Python client for CockroachDB with sync and async support

Skip to content

# Migration Examples

This page provides practical examples of using the migration features of the Manticore CockroachDB client. These examples demonstrate how to manage database schema changes effectively.

## Synchronous Migrations

This example shows how to create and apply migrations using the synchronous API:

    import os
    from manticore_cockroachdb import Database, Migration

    def sync_migration_example():
        # Connect to database
        db = Database(database="testdb", host="localhost")
        db.connect()

        try:
            # Create migration directory if it doesn't exist
            migrations_dir = "migrations"
            if not os.path.exists(migrations_dir):
                os.makedirs(migrations_dir)

            # Create migration files
            Migration.create_migration(
                "001_create_users_table",
                """
                CREATE TABLE users (
                    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    username TEXT NOT NULL UNIQUE,
                    email TEXT NOT NULL,
                    created_at TIMESTAMP NOT NULL DEFAULT NOW()
                );
                """,
                """
                DROP TABLE users;
                """,
                directory=migrations_dir
            )

            Migration.create_migration(
                "002_add_age_column",
                """
                ALTER TABLE users ADD COLUMN age INTEGER;
                """,
                """
                ALTER TABLE users DROP COLUMN age;
                """,
                directory=migrations_dir
            )

            # Load migrations
            migrations = Migration.load_migrations(migrations_dir)
            print("Loaded {} migrations".format(len(migrations)))

            # Print migration details
            for m in migrations:
                print("  Version {}: {}".format(m.version, m.description))

            # Apply all migrations
            applied = Migration.apply_migrations(db, migrations)
            print("Applied {} migrations".format(len(applied)))

            # Insert test data
            db.execute(
                "INSERT INTO users (username, email, age) VALUES (%s, %s, %s)",
                ("testuser", "test@example.com", 30)
            )

            # Query the data
            with db.cursor() as cursor:
                cursor.execute("SELECT * FROM users")
                users = cursor.fetchall()
                for user in users:
                    print("User: {}, Email: {}, Age: {}".format(
                        user["username"], user["email"], user["age"]
                    ))

            # Check current migration version
            version = Migration.get_current_version(db)
            print("Current migration version: {}".format(version))

            # Roll back the last migration
            rolled_back = Migration.rollback_migration(db, migrations[-1])
            print("Rolled back migration: {}".format(rolled_back.version))

            # Verify the age column is gone
            with db.cursor() as cursor:
                cursor.execute("SELECT * FROM users")
                users = cursor.fetchall()
                for user in users:
                    print("User after rollback: {}, Email: {}".format(
                        user["username"], user["email"]
                    ))
                    # Age column should no longer exist

        finally:
            # Clean up
            try:
                db.execute("DROP TABLE IF EXISTS migrations")
                db.execute("DROP TABLE IF EXISTS users")
            except Exception as e:
                print("Cleanup error: {}".format(e))
            db.close()

    if __name__ == "__main__":
        sync_migration_example()

## Asynchronous Migrations

This example shows how to create and apply migrations using the asynchronous API:

    import os
    import asyncio
    from manticore_cockroachdb import AsyncDatabase, AsyncMigration

    async def async_migration_example():
        # Connect to database
        db = AsyncDatabase(database="testdb", host="localhost")
        await db.connect()

        try:
            # Create migration directory if it doesn't exist
            migrations_dir = "async_migrations"
            if not os.path.exists(migrations_dir):
                os.makedirs(migrations_dir)

            # Create migration files
            await AsyncMigration.create_migration(
                "001_create_products_table",
                """
                CREATE TABLE products (
                    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    name TEXT NOT NULL,
                    price DECIMAL(10,2) NOT NULL,
                    created_at TIMESTAMP NOT NULL DEFAULT NOW()
                );
                """,
                """
                DROP TABLE products;
                """,
                directory=migrations_dir
            )

            await AsyncMigration.create_migration(
                "002_add_description_column",
                """
                ALTER TABLE products ADD COLUMN description TEXT;
                """,
                """
                ALTER TABLE products DROP COLUMN description;
                """,
                directory=migrations_dir
            )

            await AsyncMigration.create_migration(
                "003_add_category_column",
                """
                ALTER TABLE products ADD COLUMN category TEXT NOT NULL DEFAULT 'Uncategorized';
                """,
                """
                ALTER TABLE products DROP COLUMN category;
                """,
                directory=migrations_dir
            )

            # Load migrations
            migrations = await AsyncMigration.load_migrations(migrations_dir)
            print("Loaded {} migrations".format(len(migrations)))

            # Print migration details
            for m in migrations:
                print("  Version {}: {}".format(m.version, m.description))

            # Apply all migrations
            applied = await AsyncMigration.apply_migrations(db, migrations)
            print("Applied {} migrations".format(len(applied)))

            # Insert test data
            async with db.cursor() as cursor:
                await cursor.execute(
                    "INSERT INTO products (name, price, description, category) VALUES (%s, %s, %s, %s)",
                    ("Laptop", 1299.99, "High-performance laptop", "Electronics")
                )

                await cursor.execute(
                    "INSERT INTO products (name, price, description, category) VALUES (%s, %s, %s, %s)",
                    ("Headphones", 199.99, "Noise-cancelling headphones", "Audio")
                )

            # Query the data
            async with db.cursor() as cursor:
                await cursor.execute("SELECT * FROM products")
                products = await cursor.fetchall()
                for product in products:
                    print("Product: {}, Price: ${:.2f}, Category: {}, Description: {}".format(
                        product["name"], product["price"], product["category"], product["description"]
                    ))

            # Check current migration version
            version = await AsyncMigration.get_current_version(db)
            print("Current migration version: {}".format(version))

            # Roll back the last migration
            rolled_back = await AsyncMigration.rollback_migration(db, migrations[-1])
            print("Rolled back migration: {}".format(rolled_back.version))

            # Roll back another migration
            rolled_back = await AsyncMigration.rollback_migration(db, migrations[-2])
            print("Rolled back migration: {}".format(rolled_back.version))

            # Verify columns are gone
            async with db.cursor() as cursor:
                await cursor.execute("SELECT * FROM products")
                products = await cursor.fetchall()
                for product in products:
                    print("Product after rollbacks: {}, Price: ${:.2f}".format(
                        product["name"], product["price"]
                    ))
                    # Description and category columns should no longer exist

        finally:
            # Clean up
            try:
                await db.execute("DROP TABLE IF EXISTS migrations")
                await db.execute("DROP TABLE IF EXISTS products")
            except Exception as e:
                print("Cleanup error: {}".format(e))
            await db.close()

    if __name__ == "__main__":
        asyncio.run(async_migration_example())

## Migration with Transactions

This example demonstrates how to use transactions with migrations for atomic schema changes:

    import os
    from manticore_cockroachdb import Database, Migration

    def transaction_migration_example():
        # Connect to database
        db = Database(database="testdb", host="localhost")
        db.connect()

        try:
            # Create migration directory
            migrations_dir = "transaction_migrations"
            if not os.path.exists(migrations_dir):
                os.makedirs(migrations_dir)

            # Create a migration with multiple operations in a transaction
            Migration.create_migration(
                "001_create_related_tables",
                """
                -- Create categories table
                CREATE TABLE categories (
                    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    name TEXT NOT NULL UNIQUE,
                    description TEXT
                );

                -- Create items table with foreign key
                CREATE TABLE items (
                    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    name TEXT NOT NULL,
                    category_id UUID NOT NULL REFERENCES categories(id),
                    price DECIMAL(10,2) NOT NULL,
                    created_at TIMESTAMP NOT NULL DEFAULT NOW()
                );

                -- Create index for faster lookups
                CREATE INDEX items_category_idx ON items(category_id);
                """,
                """
                DROP TABLE IF EXISTS items;
                DROP TABLE IF EXISTS categories;
                """,
                directory=migrations_dir
            )

            # Load and apply migrations
            migrations = Migration.load_migrations(migrations_dir)
            applied = Migration.apply_migrations(db, migrations)
            print("Applied {} migrations".format(len(applied)))

            # Insert test data using transactions
            def insert_test_data(conn):
                with conn.cursor() as cur:
                    # Insert categories
                    cur.execute(
                        "INSERT INTO categories (name, description) VALUES (%s, %s) RETURNING id",
                        ("Electronics", "Electronic devices and gadgets")
                    )
                    electronics_id = cur.fetchone()["id"]

                    cur.execute(
                        "INSERT INTO categories (name, description) VALUES (%s, %s) RETURNING id",
                        ("Books", "Books and publications")
                    )
                    books_id = cur.fetchone()["id"]

                    # Insert items
                    items = [
                        ("Smartphone", electronics_id, 699.99),
                        ("Tablet", electronics_id, 349.99),
                        ("Novel", books_id, 14.99),
                        ("Textbook", books_id, 79.99)
                    ]

                    for name, category_id, price in items:
                        cur.execute(
                            "INSERT INTO items (name, category_id, price) VALUES (%s, %s, %s)",
                            (name, category_id, price)
                        )

                    return {
                        "electronics_id": electronics_id,
                        "books_id": books_id
                    }

            # Run in transaction
            result = db.run_in_transaction(insert_test_data)
            print("Inserted test data with category IDs: {}".format(result))

            # Query the data with a join
            with db.cursor() as cursor:
                cursor.execute("""
                    SELECT i.name as item_name, i.price, c.name as category_name
                    FROM items i
                    JOIN categories c ON i.category_id = c.id
                    ORDER BY c.name, i.price DESC
                """)

                items = cursor.fetchall()
                print("\nItems by category:")
                for item in items:
                    print("  {} ({}): ${:.2f}".format(
                        item["item_name"], item["category_name"], item["price"]
                    ))

        finally:
            # Clean up
            try:
                db.execute("DROP TABLE IF EXISTS items")
                db.execute("DROP TABLE IF EXISTS categories")
                db.execute("DROP TABLE IF EXISTS migrations")
            except Exception as e:
                print("Cleanup error: {}".format(e))
            db.close()

    if __name__ == "__main__":
        transaction_migration_example()

## Migration Status and History

This example shows how to check migration status and history:

```python import os from manticore_cockroachdb import Database, Migration

def migration_status_example(): # Connect to database db = Database(database="testdb", host="localhost") db.connect()

    try:
        # Create migration directory
        migrations_dir = "status_migrations"
        if not os.path.exists(migrations_dir):
            os.makedirs(migrations_dir)

        # Create several migrations
        for i in range(1, 6):
            version = f"{i:03d}"
            Migration.create_migration(
                f"{version}_migration_{i}",
                f"-- Migration {i} up\nCREATE TABLE IF NOT EXISTS table_{i} (id SERIAL PRIMARY KEY, name TEXT);",
                f"-- Migration {i} down\nDROP TABLE IF EXISTS table_{i};",
                directory=migrations_dir
            )

        # Load migrations
        migrations = Migration.load_migrations(migrations_dir)
        print("Loaded {} migrations".format(len(migrations)))

        # Apply only the first 3 migrations
        applied = Migration.apply_migrations(db, migrations[:3])
        print("Applied {} migrations".format(len(applied)))

        # Check current version
        version = Migration.get_current_version(db)
        print("Current migration version: {}".format(version))

        # Get migration history
        history = Migration.get_migration_history(db)
        print("\nMigration history:")
        for entry in history:
            print("  Version: {}, Applied at: {}, Description: {}".format(
                entry["version"], entry["applied_at"], entry["description"]
            ))

        # Check pending migrations
        pending = Migration.get_pending_migrations(db, migrations)
        print("\nPending migrations:")
        for m in pending:
            print("  Version: {}, Description: {}".format(m.version, m.description))

        # Apply one more migration
        applied = Migration.apply_migrations(db, [pending[0]])
        print("\nApplied 1 more migration: {}".format(applied[0].version))

        # Check status again
        version = Migration.get_current_version(db)
        print("Current migration version: {}".format(version))

        # Roll back to version 002
        target_version = "002"
        rolled_back = Migration.rollback_to_version(db, migrations, target_version)
        print("\nRolled back to version {}, migrations rolled back: {}".format(
            target_version, len(rolled_back)
        ))

        # Verify current version
        version = Migration.get_current_version(db)
        print("Current migration version after rollback: {}".format(version))

    finally:
        # Clean up
        try:
            for i in range(1, 6):
                db.execute(f"DROP TABLE IF EXISTS table_{i}")
            db.execute("DROP TABLE IF EXISTS migrations")
        except Exception as e:
            print("Cleanup error: {}".format(e))
        db.close()

if name == "main": migration_status_example()