backend/code_review_backend/issues/migrations/0007_issuelink_swap.py

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

# Generated by Django 4.1.2 on 2022-12-07 16:59

from math import ceil

from django.db import migrations

ISSUES_INSERT_SIZE = 1000


def generate_issue_links(apps, schema_editor):
    """Generate the IssueLink M2M table from issues' FK to the diff of a revision.
    Issues are migrated from the most recent to the oldest.
    The generation is done in small transactions, so the tables ar not locked for a
    long period of time (allowing to ingest new issues) and can be resumed later on.
    """
    Issue = apps.get_model("issues", "Issue")
    IssueLink = apps.get_model("issues", "IssueLink")

    # Search for the IssueLink referencing the older issue
    older_issue_link = IssueLink.objects.order_by("issue__created").first()
    issue_filters = {}
    if older_issue_link is not None:
        issue_filters = {"created__lt": older_issue_link.issue.created}

    qs = (
        # This ensures we do not handle newly created issues
        # during the migration so the order is preserved
        Issue.objects.filter(diff__isnull=False)
        .filter(**issue_filters)
        .order_by("-created", "id")
        .values("id", "diff_id", "diff__revision_id")
    )

    issues_count = qs.count()
    slices = ceil(issues_count / ISSUES_INSERT_SIZE)

    for index in range(0, slices):
        current = min((index + 1) * ISSUES_INSERT_SIZE, issues_count)
        print(
            f"[{current}/{issues_count}] Initializing Issues references on the M2M table."
        )
        issues = qs[index * ISSUES_INSERT_SIZE : (index + 1) * ISSUES_INSERT_SIZE]
        IssueLink.objects.bulk_create(
            IssueLink(
                issue_id=issue["id"],
                diff_id=issue["diff_id"],
                revision_id=issue["diff__revision_id"],
            )
            for issue in issues
        )


class Migration(migrations.Migration):
    atomic = False

    dependencies = [
        ("issues", "0006_issuelink_initial"),
    ]

    operations = [
        # Fill the M2M table
        migrations.RunPython(
            generate_issue_links,
            reverse_code=None,
            elidable=True,
            atomic=False,
        ),
        # Drop old FK
        migrations.RemoveField(
            model_name="issue",
            name="diff",
        ),
    ]