# Generated by Django 3.2.14 on 2022-07-21 22:35
from django.db import migrations


# This migration is supposed to be zero downtime safe.
# It will attempt to populate the sha256 field wherever possible. But it is expected that wherever
# older versions pulp_ansible are still running, CollectionVersions without this digest will be
# created.


def add_sha256_to_current_models(apps, schema_editor):
    """Adds the sha256 to current CollectionVersion models."""
    CollectionVersion = apps.get_model("ansible", "CollectionVersion")
    collection_versions_to_update = []
    collection_versions_on_demand = []

    for collection_version in (
        CollectionVersion.objects.prefetch_related(
            "contentartifact_set", "contentartifact_set__artifact"
        )
        .filter(sha256="")
        .only("pk", "sha256")
        .iterator(chunk_size=2000)
    ):
        content_artifact = collection_version.contentartifact_set.get()
        if content_artifact.artifact:
            collection_version.sha256 = content_artifact.artifact.sha256
            collection_versions_to_update.append(collection_version)
        else:
            collection_versions_on_demand.append(collection_version)
        if len(collection_versions_to_update) >= 1024:
            CollectionVersion.objects.bulk_update(
                collection_versions_to_update,
                [
                    "sha256",
                ],
            )
            collection_versions_to_update.clear()
    # Update remaining collection versions
    if len(collection_versions_to_update) > 0:
        CollectionVersion.objects.bulk_update(
            collection_versions_to_update,
            [
                "sha256",
            ],
        )

    # If there are on-demand collections then the next migration will fail, so error here with
    # helpful message on how to fix. No work will be performed by this migration on a second-run.
    if len(collection_versions_on_demand) > 0:
        raise Exception(
            f"On demand collections found. Please remove or upload/sync their data: "
            f"{[c.pk for c in collection_versions_on_demand]}"
        )


class Migration(migrations.Migration):

    dependencies = [
        ("ansible", "0056_collectionversion_sha256"),
    ]

    operations = [
        migrations.RunPython(add_sha256_to_current_models, migrations.RunPython.noop, elidable=True)
    ]
