Add migration for missing metadata artifacts

Author: jobselko

pulp_python/app/migrations/0019_create_missing_metadata_artifacts.py

@@ -0,0 +1,198 @@
+# Generated manually on 2025-12-15 14:00 for creating missing metadata artifacts
+
+from django.db import migrations
+
+
+def pulp_hashlib_new(name, *args, **kwargs):
+    """
+    Copied and updated (to comply with migrations) from pulpcore.
+    """
+    import hashlib as the_real_hashlib
+    from django.conf import settings
+
+    if name not in settings.ALLOWED_CONTENT_CHECKSUMS:
+        return None
+
+    return the_real_hashlib.new(name, *args, **kwargs)
+
+
+def init_and_validate(file, artifact_model, expected_digests):
+    """
+    Copied and updated (to comply with migrations) from pulpcore.
+    """
+    from django.conf import settings
+
+    digest_fields = []
+    for alg in ("sha512", "sha384", "sha256", "sha224", "sha1", "md5"):
+        if alg in settings.ALLOWED_CONTENT_CHECKSUMS:
+            digest_fields.append(alg)
+
+    if isinstance(file, str):
+        with open(file, "rb") as f:
+            hashers = {
+                n: hasher for n in digest_fields if (hasher := pulp_hashlib_new(n)) is not None
+            }
+            if not hashers:
+                return None
+
+            size = 0
+            while True:
+                chunk = f.read(1048576)  # 1 megabyte
+                if not chunk:
+                    break
+                for algorithm in hashers.values():
+                    algorithm.update(chunk)
+                size = size + len(chunk)
+    else:
+        size = file.size
+        hashers = file.hashers
+
+    for algorithm, expected_digest in expected_digests.items():
+        if algorithm not in hashers:
+            return None
+        actual_digest = hashers[algorithm].hexdigest()
+        if expected_digest != actual_digest:
+            return None
+            # todo: log/fail, user should repair sha256
+
+    attributes = {"size": size, "file": file}
+    for algorithm in digest_fields:
+        attributes[algorithm] = hashers[algorithm].hexdigest()
+
+    return artifact_model(**attributes)
+
+
+def extract_wheel_metadata(filename):
+    """
+    Extract the metadata file content from a wheel file.
+    Returns the raw metadata content as bytes or None if metadata cannot be extracted.
+    """
+    import zipfile
+
+    try:
+        with zipfile.ZipFile(filename, "r") as f:
+            for file_path in f.namelist():
+                if file_path.endswith(".dist-info/METADATA"):
+                    return f.read(file_path)
+    except (zipfile.BadZipFile, KeyError, OSError):
+        pass
+    return None
+
+
+def artifact_to_metadata_artifact(filename, artifact, md_digests, tmp_dir, artifact_model):
+    """
+    Creates artifact for metadata from the provided wheel artifact.
+    """
+    import shutil
+    import tempfile
+
+    if not filename.endswith(".whl"):
+        return None
+
+    with tempfile.NamedTemporaryFile("wb", dir=tmp_dir, suffix=filename, delete=False) as temp_file:
+        temp_wheel_path = temp_file.name
+        artifact.file.seek(0)
+        shutil.copyfileobj(artifact.file, temp_file)
+        temp_file.flush()
+
+    metadata_content = extract_wheel_metadata(temp_wheel_path)
+    if not metadata_content:
+        return None
+
+    with tempfile.NamedTemporaryFile(
+        "wb", dir=tmp_dir, suffix=".metadata", delete=False
+    ) as temp_md:
+        temp_metadata_path = temp_md.name
+        temp_md.write(metadata_content)
+        temp_md.flush()
+
+    metadata_artifact = init_and_validate(temp_metadata_path, artifact_model, md_digests)
+    if not metadata_artifact:
+        return None
+
+    try:
+        metadata_artifact.save()
+    except Exception:
+        return None
+
+    return metadata_artifact
+
+
+# todo: bulk create
+def create_missing_metadata_artifacts(apps, schema_editor):
+    """
+    Create metadata artifacts for PythonPackageContent instances that have metadata_sha256
+    but are missing the corresponding metadata artifact.
+    """
+    import tempfile
+    from django.conf import settings
+
+    PythonPackageContent = apps.get_model("python", "PythonPackageContent")
+    ContentArtifact = apps.get_model("core", "ContentArtifact")
+    Artifact = apps.get_model("core", "Artifact")
+
+    # todo: filter only wheels!
+    packages = (
+        PythonPackageContent.objects.filter(metadata_sha256__isnull=False)
+        .exclude(metadata_sha256="")
+        .prefetch_related("contentartifact_set")
+        .only("filename", "metadata_sha256")
+    )
+    # todo: only for testing, remove later
+    created_count = 0
+    skipped_count = 0
+
+    with tempfile.TemporaryDirectory(dir=settings.WORKING_DIRECTORY) as temp_dir:
+        for package in packages:
+            filename = package.filename
+            metadata_relative_path = f"{filename}.metadata"
+            content_artifacts = list(package.contentartifact_set.all())
+
+            if any(ca.relative_path == metadata_relative_path for ca in content_artifacts):
+                # Metadata artifact already exist
+                continue
+
+            main_content_artifact = next(
+                (ca for ca in content_artifacts if ca.relative_path == filename),
+                None,
+            )
+            if not main_content_artifact:
+                # Main artifact does not exist
+                skipped_count += 1
+                continue
+
+            metadata_digests = {"sha256": package.metadata_sha256}
+            metadata_artifact = artifact_to_metadata_artifact(
+                filename, main_content_artifact.artifact, metadata_digests, temp_dir, Artifact
+            )
+            if not metadata_artifact:
+                # Failed to create metadata artifact
+                skipped_count += 1
+                continue
+
+            try:
+                ContentArtifact.objects.create(
+                    artifact=metadata_artifact,
+                    content=package,
+                    relative_path=metadata_relative_path,
+                )
+                created_count += 1
+            except Exception:
+                # Failed to save metadata content artifact
+                skipped_count += 1
+
+    print(f"Created {created_count} missing metadata artifacts. Skipped {skipped_count} packages.")
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("python", "0018_packageprovenance"),
+    ]
+
+    operations = [
+        migrations.RunPython(
+            create_missing_metadata_artifacts,
+            reverse_code=migrations.RunPython.noop,
+        ),
+    ]