thrillwiki_django_no_react/rides/migrations/0003_transfer_company_data.py

# Generated by Django 5.0.7 on 2024-07-25 14:30

from django.db import migrations

def transfer_company_data(apps, schema_editor):
    Company = apps.get_model('rides', 'Company')
    Ride = apps.get_model('rides', 'Ride')
    RideModel = apps.get_model('rides', 'RideModel')

    with schema_editor.connection.cursor() as cursor:
        cursor.execute("SELECT id, name, slug, description, website, founded_year, headquarters, rides_count, coasters_count FROM manufacturers_manufacturer")
        for row in cursor.fetchall():
            company, created = Company.objects.get_or_create(
                slug=row,
                defaults={
                    'name': row,
                    'description': row,
                    'website': row,
                    'founded_date': f'{row}-01-01' if row else None,
                    'headquarters': row,
                    'rides_count': row,
                    'coasters_count': row,
                    'roles': [Company.CompanyRole.MANUFACTURER]
                }
            )
            if not created and Company.CompanyRole.MANUFACTURER not in company.roles:
                company.roles.append(Company.CompanyRole.MANUFACTURER)
                company.save()
            
            Ride.objects.filter(manufacturer_id=row).update(manufacturer_id=company.id)
            RideModel.objects.filter(manufacturer_id=row).update(manufacturer_id=company.id)

        cursor.execute("SELECT id, name, slug, description, website, founded_date, headquarters FROM designers_designer")
        for row in cursor.fetchall():
            company, created = Company.objects.get_or_create(
                slug=row,
                defaults={
                    'name': row,
                    'description': row,
                    'website': row,
                    'founded_date': row,
                    'headquarters': row,
                    'roles': [Company.CompanyRole.DESIGNER]
                }
            )
            if not created and Company.CompanyRole.DESIGNER not in company.roles:
                company.roles.append(Company.CompanyRole.DESIGNER)
                company.save()

            Ride.objects.filter(designer_id=row).update(designer_id=company.id)


class Migration(migrations.Migration):

    dependencies = [
        ('rides', '0002_ridereview_ridereviewevent_ridereview_insert_insert_and_more'),
    ]

    operations = [
        migrations.RunPython(transfer_company_data),
    ]