part II of 37e69f6 - split Album.bmanager prefetches and annotations into seperate methods, tests are still WIP

Author: tykling

src/albums/api.py

@@ -71,7 +71,7 @@ def album_create(request: HttpRequest, payload: AlbumRequestSchema) -> AlbumApiR
     # assign permissions and return response
     album.add_initial_permissions()
     # get from database so the bmanager is used
-    album = Album.bmanager.get(pk=album.pk)
+    album = Album.bmanager.prefetch_active_files_list().get(pk=album.pk)
     return 201, {"bma_response": album}
 
 
@@ -83,7 +83,7 @@ def album_create(request: HttpRequest, payload: AlbumRequestSchema) -> AlbumApiR
 )
 def album_get(request: HttpRequest, album_uuid: uuid.UUID) -> AlbumApiResponseType:
     """Return an album."""
-    album = get_object_or_404(Album.bmanager.all(), uuid=album_uuid)
+    album = get_object_or_404(Album.bmanager.annotate_memberships().prefetch_active_files_list().all(), uuid=album_uuid)
     return 200, {"bma_response": album}
 
 
@@ -95,7 +95,7 @@ def album_get(request: HttpRequest, album_uuid: uuid.UUID) -> AlbumApiResponseTy
 )
 def album_list(request: HttpRequest, filters: AlbumFilters = query) -> AlbumApiResponseType:
     """Return a list of albums."""
-    albums = Album.bmanager.all()
+    albums = Album.bmanager.annotate_memberships().prefetch_active_files_list().all()
 
     if filters.files:
         # __in is OR and we want AND, build a query for .exclude() with all file UUIDs
@@ -154,7 +154,7 @@ def album_update(
     check: bool = False,
 ) -> AlbumApiResponseType:
     """Update (PATCH) or replace (PUT) an Album."""
-    album = get_object_or_404(Album.bmanager.all(), uuid=album_uuid)
+    album = get_object_or_404(Album.bmanager.prefetch_active_files_list().all(), uuid=album_uuid)
     if not request.user.has_perm("change_album", album):
         # no permission
         return 403, {"message": "Permission denied."}

src/albums/managers.py

@@ -18,30 +18,7 @@ class AlbumManager(models.Manager):  # type: ignore[type-arg]
 
     def get_queryset(self):  # type: ignore[no-untyped-def]  # noqa: ANN201
         """Annotations and prefetches for the Album model."""
-        qs = super().get_queryset()
-
-        # queries to count past, present, and future memberships of each album
-        active_memberships = Count("memberships", filter=Q(memberships__period__contains=timezone.now()), distinct=True)
-        historic_memberships = Count("memberships", filter=Q(memberships__period__endswith__lt=timezone.now()))
-        future_memberships = Count("memberships", filter=Q(memberships__period__startswith__gt=timezone.now()))
-
-        return (
-            qs.annotate(
-                active_memberships=active_memberships,
-                historic_memberships=historic_memberships,
-                future_memberships=future_memberships,
-            )
-            .select_related("owner")
-            .prefetch_related("user_permissions__user")
-            .prefetch_related("user_permissions__permission")
-            .prefetch_related("group_permissions__group")
-            .prefetch_related("group_permissions__permission")
-            .prefetch_related("hits")
-            .annotate(hitcount=Count("hits", distinct=True))
-            .prefetch_active_files_list(recursive=True)
-            # ordering from Album META gets lost for some reason :(
-        ).order_by("created_at")
-
+        return super().get_queryset()
 
 class AlbumQuerySet(models.QuerySet):  # type: ignore[type-arg]
     """Custom queryset for album operations."""
@@ -76,3 +53,30 @@ def prefetch_active_files_list(self, *, recursive: bool = True) -> models.QueryS
                 to_attr="active_files_list",
             ),
         )
+
+    def prefetch_user_permissions(self) -> models.QuerySet["Album"]:
+        """Prefetch user permissions."""
+        return self.prefetch_related("user_permissions__user").prefetch_related("user_permissions__permission")
+
+    def prefetch_group_permissions(self) -> models.QuerySet["Album"]:
+        """Prefetch group permissions."""
+        return self.prefetch_related("group_permissions__group").prefetch_related("group_permissions__permission")
+
+    def annotate_hitcount(self) -> models.QuerySet["Album"]:
+        """Annotate hitcounts for the qs."""
+        return self.annotate(hitcount=Count("hits", distinct=True))
+
+    def annotate_memberships(self) -> models.QuerySet["Album"]:
+        """Annotate membership counts."""
+        active_memberships = Count("memberships", filter=Q(memberships__period__contains=timezone.now()), distinct=True)
+        historic_memberships = Count("memberships", filter=Q(memberships__period__endswith__lt=timezone.now()))
+        future_memberships = Count("memberships", filter=Q(memberships__period__startswith__gt=timezone.now()))
+        return self.annotate(
+            active_memberships=active_memberships,
+            historic_memberships=historic_memberships,
+            future_memberships=future_memberships,
+        )
+
+    def select_owner(self) -> models.QuerySet["Album"]:
+        """Get owner with select_related()."""
+        return self.select_related("owner")

src/albums/views.py

@@ -48,7 +48,7 @@ class AlbumListView(SingleTableMixin, FilterView):
 
     def get_queryset(self) -> models.QuerySet[Album]:
         """Use bmanager to get rich album objects."""
-        return Album.bmanager.all()
+        return Album.bmanager.annotate_memberships().annotate_hitcount().all()
 
 
 class AlbumDetailView(SingleTableMixin, FilterView):
@@ -64,7 +64,7 @@ def get_template_names(self) -> list[str]:
 
     def get_object(self, queryset: models.QuerySet[Album] | None = None) -> Album:
         """Use the manager so the album object has prefetched active_files."""
-        album = Album.bmanager.get(pk=self.kwargs["album_uuid"])
+        album = Album.bmanager.prefetch_active_files_list().get(pk=self.kwargs["album_uuid"])
         # count the hit
         count_hit(self.request, album)
         return album  # type: ignore[no-any-return]

src/files/tables.py

@@ -43,7 +43,7 @@ class FileTable(tables.Table):
 
     # only show on 5xl and up
     tags = BPColumn(bp="5xl", verbose_name="Tags")
-    hitcount = BPColumn(bp="5xl", verbose_name="Hits")
+    hitcount = BPColumn(bp="5xl", verbose_name="Hits", orderable=False)
 
     # only show on 6xl and up
     jobs = BPColumn(bp="6xl", verbose_name="Jobs", order_by="jobs_finished")

src/jobs/tables.py

@@ -19,7 +19,6 @@ class JobTable(tables.Table):
     basefile = tables.Column(verbose_name="File", linkify=True)
     job_type = tables.Column(verbose_name="Job Type", orderable=False)
     result_url = tables.Column(verbose_name="Result url", orderable=False)
-
     # show only at xxl and up
     width = BPColumn(bp="xxl", orderable=False)
     height = BPColumn(bp="xxl", orderable=False)