# Generated by Django 3.2.13 on 2022-06-09 11:29

from django.db import migrations, models


def add_source_to_existing_companies_entries(apps, schema_editor):
    # add "source": "legacy_datahub" to all items in companies field
    Barrier = apps.get_model("barriers", "Barrier")
    for barrier in Barrier.objects.filter(companies__isnull=False):
        new_companies = [
            {
                **company,
                "source": "legacy_datahub",
            }
            for company in barrier.companies
        ]

        barrier.companies = new_companies
        barrier.save()


def remove_source_from_existing_companies_entries(apps, schema_editor):
    # remove "source": "legacy_datahub" from all items in companies field
    Barrier = apps.get_model("barriers", "Barrier")
    for barrier in Barrier.objects.filter(companies__isnull=False):
        new_companies = []
        for company in barrier.companies:
            del company["source"]
            new_companies.append(company)

        barrier.companies = new_companies
        barrier.save()


class Migration(migrations.Migration):

    dependencies = [
        ("barriers", "0127_auto_20220516_1058"),
    ]

    operations = [
        migrations.RunPython(
            add_source_to_existing_companies_entries,
            remove_source_from_existing_companies_entries,
        )
    ]
