Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Adding annotations to all querysets with a custom QuerySet as Manager

I'm not sure how to put words on this precisely, but what I'm trying to do is use a custom QuerySet for a model, but I want to annotate the results of any queryset. So whether filters were applied, or it was a direct get of an object (which I think is still a filter() call in the end), I want the annotations to be applied so I can access them in the Model objects I get back.

My current solution is using a custom QuerySet + custom Manager, and the custom Manager passes custom calls off to the appropriate call on the custom Queryset. Of course, this has the unfortunate need to match methods between the manager and queryset for any custom method I want.

I know that CustomQuerySet.as_manager() exists to avoid this duplication, but if I use that, then I lose the ability to override get_queryset() so my annotations can be applied in all cases of evaluating a queryset.

Is there a cleaner way to do this so I am still left with the ability to chain calls and have my annotations apply when the queryset is finally evaluated?

Ideally I'd like to keep using standard calls like: Model.objects.most_liked().near() and Model.objects.get(id=model_id)

As opposed to having to do something like add an add_annotations() method to my custom QuerySet, thereby needing to always call that method anytime I get Model objects (ie. I don't want Model.objects.get(id=id).add_annotations() and Model.objects.near().add_annotations() and so on)

but not have to repeat the methods between manager and queryset? Ideally without hacking into private methods or anything like that. I briefly toyed with overriding __iter__ on the CustomQuerySet class to add the annotations then, but it didn't seem right to do that.

Sample setup of model, queryset, and manager below. These of course aren't my real class names, just fillers for the question. Hopefully I've included enough to make my question understandable :)

class CustomQuerySet(models.QuerySet):
    def most_liked(self):
        return self.filter(...)

    def near(self): 
        return self.filter(...)


class CustomManager(models.Manager):
    def get_queryset(self):
        return CustomQuerySet(self.model, using=self._db)\
            .annotate(...)\
            .annotate(...)

    def near(self, latitude, longitude, radius, max_results=100):
        return self.get_queryset().near()

    def most_liked(self):
        return self.get_queryset().most_liked()

class Model(models.Model):
    objects = CustomManager()
like image 230
Shane Avatar asked Feb 28 '17 22:02

Shane


1 Answers

You should use the from_queryset method - it allows you to create a Manager class which you can further customise before applying to your model. For example:

class CustomQuerySet(models.QuerySet):
    def most_liked(self):
        return self.filter(...)


class CustomManager(models.Manager.from_queryset(CustomQuerySet)):
    def get_queryset(self):
        return super(CustomManager, self).get_queryset() \
            .annotate(...)


class Model(models.Model):
    objects = CustomManager()
like image 53
Greg Avatar answered Oct 21 '22 03:10

Greg