<gh_stars>1000+
# Generated by Django 3.2.12 on 2022-04-25 08:31

from django.db import migrations
from django.db.models import Exists, OuterRef

BATCH_SIZE = 5000


def rewrite_order_events_parameters(apps, _schema_editor):
    """Rewrite order events that have old `line_pk` value to `uuid` value."""
    Order = apps.get_model("order", "Order")
    OrderLine = apps.get_model("order", "OrderLine")
    OrderEvent = apps.get_model("order", "OrderEvent")

    queryset = OrderEvent.objects.filter(parameters__lines__isnull=False)
    for batch_pks in queryset_in_batches(queryset):
        events = OrderEvent.objects.filter(pk__in=batch_pks)
        orders = Order.objects.filter(Exists(events.filter(order_id=OuterRef("id"))))
        order_line_old_id_to_new_mapping = {
            line_ids["old_id"]: str(line_ids["id"])
            for line_ids in OrderLine.objects.filter(order__in=orders).values(
                "old_id", "id"
            )
        }
        for event in events:
            lines = event.parameters["lines"]
            for line_data in lines:
                line_data["line_pk"] = order_line_old_id_to_new_mapping.get(
                    line_data["line_pk"]
                )
        OrderEvent.objects.bulk_update(events, ["parameters"])


def queryset_in_batches(queryset):
    """Slice a queryset into batches.

    Input queryset should be sorted be pk.
    """
    start_pk = 0

    while True:
        qs = queryset.order_by("pk").filter(pk__gt=start_pk)[:BATCH_SIZE]
        pks = list(qs.values_list("pk", flat=True))

        if not pks:
            break

        yield pks

        start_pk = pks[-1]


class Migration(migrations.Migration):

    dependencies = [
        ("order", "0144_rewrite_fulfillmentline_orderline_relation"),
    ]

    operations = [
        migrations.RunPython(
            rewrite_order_events_parameters,
            migrations.RunPython.noop,
        ),
    ]
