quantum-ai-v2/agent_base/management/commands/reset_database.py

from django.core.management.base import BaseCommand
from django.core.management import call_command
from django.db import connection, transaction
from django.conf import settings
import os
import shutil


class Command(BaseCommand):
    help = 'Reset database and migrations for clean development/deployment'
    
    def add_arguments(self, parser):
        parser.add_argument(
            '--action',
            choices=['migrations', 'database', 'full'],
            default='full',
            help='What to reset: migrations, database, or full (both)',
        )
        parser.add_argument(
            '--confirm',
            action='store_true',
            help='Confirm the destructive action',
        )
        parser.add_argument(
            '--keep-superuser',
            action='store_true',
            help='Keep existing superuser data during database reset',
        )
    
    def handle(self, *args, **options):
        action = options['action']
        confirm = options['confirm']
        keep_superuser = options['keep_superuser']
        
        if not confirm:
            self.stdout.write(
                self.style.WARNING(
                    "⚠️  This is a destructive operation! Add --confirm to proceed."
                )
            )
            self.stdout.write("This will:")
            if action in ['migrations', 'full']:
                self.stdout.write("  - Delete all migration files")
            if action in ['database', 'full']:
                self.stdout.write("  - Drop all database tables")
                self.stdout.write("  - Recreate fresh database")
            return
        
        if action in ['migrations', 'full']:
            self.reset_migrations()
        
        if action in ['database', 'full']:
            self.reset_database(keep_superuser)
        
        if action == 'full':
            self.create_fresh_migrations()
            self.run_migrations()
            if not keep_superuser:
                self.create_initial_data()
    
    def reset_migrations(self):
        """Delete all migration files except __init__.py"""
        self.stdout.write("🗑️  Deleting migration files...")
        
        apps_with_migrations = [
            'agent_base',
            'authentication', 
            'core',
            'wallet',
            'weather_reporter',
            'data_analyzer',
            'job_posting_generator',
            'social_ads_generator',
        ]
        
        for app in apps_with_migrations:
            migrations_dir = f"{app}/migrations"
            if os.path.exists(migrations_dir):
                # Keep __init__.py but delete all other migration files
                for file in os.listdir(migrations_dir):
                    if file.endswith('.py') and file != '__init__.py':
                        file_path = os.path.join(migrations_dir, file)
                        os.remove(file_path)
                        self.stdout.write(f"  Deleted: {file_path}")
        
        self.stdout.write(self.style.SUCCESS("✅ Migration files deleted"))
    
    def reset_database(self, keep_superuser=False):
        """Drop all tables and recreate database"""
        self.stdout.write("🗑️  Resetting database...")
        
        # Backup superuser if requested
        superuser_data = None
        if keep_superuser:
            superuser_data = self.backup_superuser()
        
        # Get database engine
        db_config = settings.DATABASES['default']
        engine = db_config['ENGINE']
        
        if 'sqlite' in engine:
            # For SQLite, just delete the file
            db_file = db_config['NAME']
            if os.path.exists(db_file):
                os.remove(db_file)
                self.stdout.write(f"  Deleted SQLite file: {db_file}")
        
        elif 'postgresql' in engine:
            # For PostgreSQL, drop all tables
            self.drop_all_postgresql_tables()
        
        else:
            self.stdout.write(
                self.style.ERROR(f"Unsupported database engine: {engine}")
            )
            return
        
        self.stdout.write(self.style.SUCCESS("✅ Database reset"))
        
        # Restore superuser if backed up
        if superuser_data:
            self.restore_superuser(superuser_data)
    
    def drop_all_postgresql_tables(self):
        """Drop all tables in PostgreSQL database"""
        with connection.cursor() as cursor:
            # Get all table names
            cursor.execute("""
                SELECT tablename FROM pg_tables 
                WHERE schemaname = 'public'
            """)
            tables = [row[0] for row in cursor.fetchall()]
            
            if tables:
                # Drop all tables with CASCADE
                tables_str = ', '.join(f'"{table}"' for table in tables)
                cursor.execute(f'DROP TABLE IF EXISTS {tables_str} CASCADE')
                self.stdout.write(f"  Dropped {len(tables)} PostgreSQL tables")
    
    def backup_superuser(self):
        """Backup superuser data before reset"""
        try:
            from django.contrib.auth import get_user_model
            User = get_user_model()
            
            superuser = User.objects.filter(is_superuser=True).first()
            if superuser:
                return {
                    'username': superuser.username,
                    'email': superuser.email,
                    'first_name': superuser.first_name,
                    'last_name': superuser.last_name,
                }
        except Exception:
            pass
        return None
    
    def restore_superuser(self, superuser_data):
        """Restore superuser after reset"""
        if superuser_data:
            self.stdout.write("🔑 Restoring superuser...")
            call_command(
                'create_user',
                superuser_data['email'],
                'admin123',  # Default password
                '--superuser',
                '--username', superuser_data['username'],
                '--first-name', superuser_data['first_name'],
                '--last-name', superuser_data['last_name'],
            )
    
    def create_fresh_migrations(self):
        """Create new migration files"""
        self.stdout.write("📝 Creating fresh migrations...")
        call_command('makemigrations')
        self.stdout.write(self.style.SUCCESS("✅ Fresh migrations created"))
    
    def run_migrations(self):
        """Apply all migrations"""
        self.stdout.write("🔄 Running migrations...")
        call_command('migrate')
        self.stdout.write(self.style.SUCCESS("✅ Migrations applied"))
    
    def create_initial_data(self):
        """Create initial data (agents and admin user)"""
        self.stdout.write("👤 Creating initial data...")
        call_command('populate_agents', '--create-admin')
        self.stdout.write(self.style.SUCCESS("✅ Initial data created"))