🗄️ Database Schema Migration Agent

#!/usr/bin/env python3
"""
Database Schema Migration Agent

This agent generates database migration scripts for schema changes including
adding/removing columns, creating indices, and data transformations. It ensures
zero-downtime migrations with rollback procedures.

Usage:
    python 06_db_migration.py "Add NOT NULL email column to users table"
    python 06_db_migration.py "Rename user_id to customer_id"
"""

import sys
from anthropic import Anthropic

def main():
    if len(sys.argv) < 2:
        print("Usage: python 06_db_migration.py <migration_description>")
        sys.exit(1)

    target = " ".join(sys.argv[1:])

    client = Anthropic()

    system_prompt = """You are a database migration expert. Your job is to:
1. Generate safe, zero-downtime migrations
2. Create backfill procedures for data transformation
3. Implement proper indexing strategies
4. Provide rollback procedures
5. Identify locking issues and mitigation
6. Support multiple database systems (PostgreSQL, MySQL, SQLite)

Always prioritize data integrity and backward compatibility."""

    task = f"""Generate a complete database migration for:

{target}

Provide:
1. Pre-migration checks and validation
2. Migration script with:
   - Initial state
   - New state
   - Data transformation (if needed)
   - Index creation/removal
3. Zero-downtime strategy (if applicable)
4. Rollback procedure with reverse migrations
5. Performance impact analysis
6. Testing strategy and test SQL queries
7. Deployment checklist
8. Monitoring recommendations

Include comments explaining each step."""

    print("🗄️  Starting Database Migration Agent...\n")
    print("=" * 60)

    try:
        with client.messages.stream(
            model="claude-opus-4-7",
            max_tokens=3000,
            system=system_prompt,
            messages=[{"role": "user", "content": task}],
        ) as stream:
            for text in stream.text_stream:
                print(text, end="", flush=True)

        print("\n" + "=" * 60)
        print("\n✅ Migration generation complete!")

    except Exception as e:
        print(f"\n❌ Error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()