Annotated paginator page

12 Nov 2022

First things first. Here’s the code. Enjoy!

from django.contrib import admin
from django.core.paginator import Paginator, Page
from django.db.models import Count, Exists, OuterRef

from myapp.models import MyModel, OtherModel


class AnnotatedPaginator(Paginator):
    """
    Apply a dictionary of annotations to the page for the paginator.
    """
    
    def __init__(self, *args, annotations=None, **kwargs):
        self.annotations = annotations
        super().__init__(*args, **kwargs)

    def _get_page(self, object_list, *args, **kwargs):
        """
        Return an instance of a single page.
        This will change the object_list into an actual list.
        It will make an additional query to the database to look up
        the values to be manually set on the object.
        """
        objects = list(object_list)
        if objects and self.annotations:
            # Make another query for this model type to gather the annotated fields.
            annotated_queryset = (
                objects[0]
                ._meta.model.objects.filter(id__in=[obj.id for obj in objects])
                .annotate(**self.annotations)
            )
            # Create a map to associate the original objects to the annotated values.
            annotated_maps = {
                annotated_map.pop("id"): annotated_map
                for annotated_map in annotated_queryset.values(
                    "id", *self.annotations.keys()
                )
            }
            # Associated the annotated values to the original objects.
            for obj in objects:
                for key, value in annotated_maps[obj.id].items():
                    setattr(obj, key, value)
        return Page(objects, *args, **kwargs)


class AdminAnnotatedPageMixin:
    """
    Extend the ModelAdmin functionality to utilize AnnotatedPaginator
    """
    
    paginator = AnnotatedPaginator
    page_annotations = None

    def get_paginator(
        self, request, queryset, per_page, orphans=0, allow_empty_first_page=True
    ):
        return self.paginator(
            queryset,
            per_page,
            orphans,
            allow_empty_first_page,
            annotations=self.page_annotations,
        )


@admin.site.register(MyModel)
class MyModelAdmin(AdminAnnotatedPageMixin, admin.ModelAdmin):
    page_annotations = {
        "related_field_count": Count("related_field"),
        "special_case_exists": Exists(
            OtherModel.objects.filter(relation_id=OuterRef("id"))[:1]
        ),
    }
    list_display = ['id', 'related_field_count', 'special_case_exists']
    
    @admin.display(description="Related Count")
    def related_field_count(self, obj):
        return obj.related_field_count
    
    @admin.display(description="Is Special", boolean=True)
    def special_case_exists(self, obj):
        return obj.special_case_exists

If you’re curious about what the above does here’s my explanation.

I ran across the use-case recently where I wanted to show additional context for a model’s admin list view. This model had about a million rows and one of the related models’ tables was in the 10 million to 100 million row range. So not small, but not massive. I was running into a problem where the two annotations were being nested resulting in a query that should take 100ms taking 10-30s!

I couldn’t identify a good index or database schema solution, so I decided to try to solve this by only applying the annotation to the query that selects the data. Ignoring the hacky solution I attempted with ModelAdmin.get_changelist_instance, I ended up subclassing the Paginator class to ioverride the _get_page hook so that I could apply the annotation exactly where I wanted.

In the above code, AnnotatedPaginator accepts an annotations dictionary that will be passed into a .annotate() function call. Those annotations are applied within the _get_page() function.

The approach that my solution takes is to convert the sliced data into a list, then perform another database query to look up the desired information only for the instances in that slice. *It is possible to call annotate on a sliced QuerySet, but there’s no restriction that object_list must be a QuerySet. I decided to implement for robustness rather than performance. If I were to package this up, I would do the inspection on object_list and implement for optimization.

At this point we have our paginator class that can support the annotation look up, so now we need to integrate it with Django’s admin site. To help with this, I created a mixin class that overrides the get_paginator method to supply the annotations parameter. This class allows you to specify the annotations as a class member in your own ModelAdmin class following Django’s pattern.

This means you end up with a pretty compact definition:

@admin.site.register(MyModel)
class MyModelAdmin(AdminAnnotatedPageMixin, admin.ModelAdmin):
    page_annotations = {
        "related_field_count": Count("related_field"),
        "special_case_exists": Exists(
            OtherModel.objects.filter(relation_id=OuterRef("id"))[:1]
        ),
    }
    list_display = ['id', 'related_field_count', 'special_case_exists']
    
    @admin.display(description="Related Count")
    def related_field_count(self, obj):
        return obj.related_field_count
    
    @admin.display(description="Is Special", boolean=True)
    def special_case_exists(self, obj):
        return obj.special_case_exists

Note that you still must add specific functions to expose those annotated values on the list view. Additionally, you won’t be able to sort on these columns since they are only being calculated for the specific page of data.