Automate Database Migrations with AI Agents in 25 Minutes

pip install langchain-anthropic sqlalchemy alembic --break-system-packages

# migration_agent.py
from langchain_anthropic import ChatAnthropic
from langchain.agents import Tool, AgentExecutor, create_react_agent
from langchain.prompts import PromptTemplate
from sqlalchemy import create_engine, inspect
import json

class SchemaMigrationAgent:
    def __init__(self, db_url: str, api_key: str):
        self.engine = create_engine(db_url)
        self.llm = ChatAnthropic(
            model="claude-sonnet-4-20250514",
            anthropic_api_key=api_key,
            temperature=0  # Deterministic for schema changes
        )
        
    def get_current_schema(self, table_name: str) -> dict:
        """Extract existing table structure"""
        inspector = inspect(self.engine)
        
        return {
            "columns": [
                {
                    "name": col["name"],
                    "type": str(col["type"]),
                    "nullable": col["nullable"],
                    "default": col.get("default")
                }
                for col in inspector.get_columns(table_name)
            ],
            "indexes": inspector.get_indexes(table_name),
            "foreign_keys": inspector.get_foreign_keys(table_name),
            "primary_key": inspector.get_pk_constraint(table_name)
        }
    
    def analyze_impact(self, table_name: str, proposed_changes: str) -> dict:
        """AI analyzes breaking change risks"""
        current_schema = self.get_current_schema(table_name)
        
        prompt = f"""You are a database migration expert. Analyze this proposed change:

Current schema for {table_name}:
{json.dumps(current_schema, indent=2)}

Proposed changes:
{proposed_changes}

Identify:
1. Breaking changes (removed columns, type changes, constraint additions)
2. Performance impacts (table locks, index rebuilds)
3. Required data migrations
4. Rollback complexity (1-5 scale)

Respond in JSON:
{{
  "breaking_changes": ["list of issues"],
  "performance_impact": "low|medium|high",
  "lock_duration_estimate": "seconds",
  "requires_data_migration": true|false,
  "rollback_complexity": 1-5,
  "safe_to_deploy": true|false,
  "recommended_approach": "step-by-step migration strategy"
}}"""

        response = self.llm.invoke(prompt)
        # Parse JSON from response
        return json.loads(response.content)

def generate_migration(self, analysis: dict, table_name: str, changes: str) -> str:
    """AI generates migration with safety checks"""
    
    if not analysis["safe_to_deploy"]:
        # Agent proposes multi-step migration for breaking changes
        prompt = f"""Generate a SAFE migration for this breaking change:

Table: {table_name}
Changes: {changes}
Issues: {analysis['breaking_changes']}

Requirements:
- Use shadow tables for destructive changes
- Add NOT NULL constraints in separate transaction
- Create indexes CONCURRENTLY (Postgres)
- Include data migration if needed
- Provide rollback script

Format as Alembic migration:"""
    else:
        prompt = f"""Generate migration SQL for:

Table: {table_name}
Changes: {changes}

Use best practices:
- IF NOT EXISTS for additive changes
- Comments explaining each step
- Performance-optimized for tables >1M rows"""

    response = self.llm.invoke(prompt)
    return response.content

# Usage example
agent = SchemaMigrationAgent(
    db_url="postgresql://user:pass@localhost/mydb",
    api_key="your_anthropic_key"
)

# Propose adding a column
proposed = "Add email_verified BOOLEAN NOT NULL DEFAULT false to users table"
analysis = agent.analyze_impact("users", proposed)

if analysis["safe_to_deploy"]:
    migration_sql = agent.generate_migration(analysis, "users", proposed)
    print(migration_sql)
else:
    print(f"⚠️  Breaking changes detected: {analysis['breaking_changes']}")
    print(f"Recommended: {analysis['recommended_approach']}")

def validate_migration(self, migration_sql: str) -> dict:
    """Test migration on database snapshot"""
    
    # Create temporary test database
    test_engine = create_engine("postgresql://localhost/test_migrations")
    
    # Clone production schema (simplified)
    with self.engine.connect() as prod, test_engine.connect() as test:
        # Copy schema only, not data (faster)
        prod.execute("pg_dump --schema-only | psql test_migrations")
        
        try:
            # Run migration
            test.execute(migration_sql)
            
            # AI validates results
            inspector = inspect(test_engine)
            new_schema = self.get_current_schema("users")
            
            validation_prompt = f"""Did this migration achieve the goal?

Before: {self.original_schema}
After: {new_schema}
Goal: {self.proposed_changes}

Check:
- All columns exist as specified
- Constraints applied correctly
- No unintended side effects

Respond: {{"valid": true/false, "issues": []}}"""
            
            result = self.llm.invoke(validation_prompt)
            return json.loads(result.content)
            
        except Exception as e:
            return {"valid": False, "issues": [str(e)]}
        finally:
            test_engine.dispose()

import time
from contextlib import contextmanager

@contextmanager
def monitored_migration(self, table_name: str):
    """Deploy with automatic rollback on failure"""
    
    connection = self.engine.connect()
    transaction = connection.begin()
    
    # Store rollback point
    connection.execute(f"SAVEPOINT before_migration")
    
    # Monitor query performance during migration
    start_time = time.time()
    initial_locks = connection.execute(
        "SELECT count(*) FROM pg_locks WHERE relation::regclass::text = %s",
        (table_name,)
    ).scalar()
    
    try:
        yield connection
        
        # AI validates post-migration state
        duration = time.time() - start_time
        if duration > 30:  # Took too long
            raise Exception(f"Migration exceeded time budget: {duration}s")
        
        transaction.commit()
        print(f"✓ Migration completed in {duration:.2f}s")
        
    except Exception as e:
        print(f"✗ Migration failed: {e}")
        connection.execute("ROLLBACK TO SAVEPOINT before_migration")
        transaction.rollback()
        
        # AI generates incident report
        report = self.llm.invoke(f"""Migration failed: {e}
        
Generate post-mortem:
- What went wrong
- How to fix manually
- Prevent future occurrence""")
        
        print(report.content)
        raise
    finally:
        connection.close()

# Deploy with monitoring
with agent.monitored_migration("users") as conn:
    conn.execute(migration_sql)

# Run against staging database
python migration_agent.py --db staging --dry-run

# Expected output:
# ✓ Schema analyzed: 0 breaking changes
# ✓ Migration generated: 15 lines SQL
# ✓ Validation passed on test database
# ✓ Safe to deploy to production

# .github/workflows/auto-migrate.yml
name: AI-Powered Migrations

on:
  pull_request:
    paths:
      - 'migrations/**'

jobs:
  validate:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      
      - name: Analyze Migration
        env:
          ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_KEY }}
          DATABASE_URL: ${{ secrets.STAGING_DB }}
        run: |
          python migration_agent.py analyze \
            --migration migrations/latest.sql \
            --output analysis.json
      
      - name: Comment on PR
        uses: actions/github-script@v7
        with:
          script: |
            const analysis = require('./analysis.json');
            github.rest.issues.createComment({
              issue_number: context.issue.number,
              body: `## Migration Analysis
              
**Breaking Changes:** ${analysis.breaking_changes.length}
**Performance Impact:** ${analysis.performance_impact}
**Safe to Deploy:** ${analysis.safe_to_deploy ? '✅' : '⚠️ Review Required'}

${analysis.recommended_approach}`
            });

class MultiServiceMigrationAgent(SchemaMigrationAgent):
    def check_cross_service_impact(self, table: str, changes: str) -> dict:
        """Scan all services for queries using this table"""
        
        # AI reads application code
        affected_services = []
        for service in self.scan_codebase():
            queries = self.extract_sql_queries(service)
            
            impact_check = self.llm.invoke(f"""
Does this migration break these queries?

Migration: {changes}
Queries from {service['name']}:
{queries}

Respond: {{"breaks": true/false, "affected_queries": []}}""")
            
            result = json.loads(impact_check.content)
            if result["breaks"]:
                affected_services.append({
                    "service": service["name"],
                    "queries": result["affected_queries"]
                })
        
        return {"affected_services": affected_services}