Update Package Set creation migration code #95

Signed-off-by: Jono Yang <[email protected]>

packagedb/migrations/0070_auto_20230706_0045.py

@@ -1,35 +1,33 @@
 # Generated by Django 4.1.2 on 2023-07-06 00:45
 
 from django.db import migrations
+from django.db.models import Count
 
 
 def create_package_sets_from_existing(apps, schema_editor):
     Package = apps.get_model("packagedb", "Package")
     PackageSet = apps.get_model("packagedb", "PackageSet")
 
-    # We first create PackageSets from Packages with package_set
-    package_set_values = Package.objects.filter(
-        package_set__isnull=False
-    ).values(
-        "package_set"
-    ).order_by(
-        "package_set"
-    ).distinct(
-        "package_set"
-    ).iterator(
-        chunk_size=5000
+    # We first create PackageSets from package_sets that contain multiple
+    # Packages
+    multi_package_set_values = (
+        Package.objects.filter(package_set__isnull=False)
+        .values("package_set")
+        .annotate(package_set_count=Count("package_set"))
+        .filter(package_set_count__gt=1)
+        .iterator(chunk_size=5000)
     )
 
     uncreated_package_sets = []
-    for i, package_set_value in enumerate(package_set_values):
+    for i, package_set_annotation in enumerate(multi_package_set_values):
         if not i % 5000:
             PackageSet.objects.bulk_create(
                 uncreated_package_sets
             )
             uncreated_package_sets = []
 
         package_set = PackageSet(
-            uuid=package_set_value["package_set"],
+            uuid=package_set_annotation["package_set"],
         )
         uncreated_package_sets.append(package_set)
 
@@ -38,17 +36,18 @@ def create_package_sets_from_existing(apps, schema_editor):
             uncreated_package_sets
         )
 
-    # We then associate those packages to the newly created PackageSet
-    # TODO: is there a way to do a bulk add for many to many relationships?
-    packages = Package.objects.filter(
-        package_set__isnull=False
-    ).iterator(
-        chunk_size=5000
+    multi_package_set_values = (
+        Package.objects.filter(package_set__isnull=False)
+        .values("package_set")
+        .annotate(package_set_count=Count("package_set"))
+        .filter(package_set_count__gt=1)
+        .iterator(chunk_size=5000)
     )
+    # Associate those packages to the new package_set
+    multi_package_sets = [m['package_set'] for m in multi_package_set_values]
+    packages = Package.objects.filter(package_set__in=multi_package_sets)
     for package in packages:
         package_set = PackageSet.objects.get(uuid=package.package_set)
-        if not package_set:
-            continue
         package_set.packages.add(package)
 
 
@@ -62,7 +61,6 @@ def create_maven_package_sets(apps, schema_editor):
         package_set__isnull=True,
         type="maven"
     ).order_by(
-        "type",
         "namespace",
         "name",
         "version",
@@ -72,9 +70,10 @@ def create_maven_package_sets(apps, schema_editor):
         chunk_size=5000
     )
 
+    prev_namespace = None
     prev_name = None
     prev_version = None
-    prev_package_set = None
+    prev_package = None
     unupdated_packages = []
     for package in maven_packages_without_package_set:
         if "source" in package.qualifiers:
@@ -84,29 +83,20 @@ def create_maven_package_sets(apps, schema_editor):
         package.package_content = package_content
         unupdated_packages.append(package)
 
-        # We first start iterating through the set or we encounter a package
-        # that is not part of the same group as the package before this
         if (
-            package.name != prev_name
-            or (
-                package.name == prev_name
-                and package.version != prev_version
-            )
+            package.namespace == prev_namespace
+            and package.name == prev_name
+            and package.version == prev_version
         ):
             package_set = PackageSet()
             package_set.save()
             package_set.packages.add(package)
+            package_set.packages.add(prev_package)
 
-            prev_name = package.name
-            prev_version = package.version
-            prev_package_set = package_set
-
-        # We are iterating past the first package and we find a related package
-        elif (
-            package.name == prev_name
-            and package.version == prev_version
-        ):
-            prev_package_set.packages.add(package)
+        prev_namespace = package.namespace
+        prev_name = package.name
+        prev_version = package.version
+        prev_package = package
 
     if unupdated_packages:
         Package.objects.bulk_update(