Annotate a Foreign Key field in Django

Question

Let's say I have three Models:

class Product(models.Model):
    name = models.CharField(max_length=255)

class Plan(models.Model):
    products = models.ManyToManyField(
        Product, 
    )

class User(models.Model):
    plan = models.ForeignKey(Plan, on_delete=models.CASCADE)

I want to make a query to get a list of Users, prefetch their plan, and annotate their plan with the count of Products. I'm trying something like this:

plan_queryset = Plan.objects.all().annotate(num_product=Count('products'));

queryset = Plan.user_set \
   .prefetch_related(
       Prefetch('plan', 
          queryset=plan_queryset
       )
   )
serializer = UserSerializer(queryset, many=True)

But unfortunately, the plan queryset doesn't execute. Passing the queryset into a serializer I get an error saying "num_product" is not defined.

What's the best way to go about annotating a foreign key field without making a query for each object? (n + 1 problem).

UPDATE:

Example Serializers:

class PlanSerializer(serializers.ModelSerializer):
    num_product=serializers.IntegerField()

    class Meta:
        model = Plan
        fields = ('num_product',)

class UserSerializer(serializers.ModelSerializer):
    plan = PlanSerializer(read_only=True)

    class Meta:
         model = User
         fields = ('plan',)

Traceback:

Traceback (most recent call last):
  File "/usr/local/lib/python3.6/site-packages/rest_framework/fields.py", line 441, in get_attribute
    return get_attribute(instance, self.source_attrs)
  File "/usr/local/lib/python3.6/site-packages/rest_framework/fields.py", line 100, in get_attribute
    instance = getattr(instance, attr)
AttributeError: 'Plan' object has no attribute 'num_product'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/local/lib/python3.6/site-packages/django/core/handlers/exception.py", line 35, in inner
    response = get_response(request)
  File "/usr/local/lib/python3.6/site-packages/django/core/handlers/base.py", line 128, in _get_response
    response = self.process_exception_by_middleware(e, request)
  File "/usr/local/lib/python3.6/site-packages/django/core/handlers/base.py", line 126, in _get_response
    response = wrapped_callback(request, *callback_args, **callback_kwargs)
  File "/usr/local/lib/python3.6/site-packages/django/views/decorators/csrf.py", line 54, in wrapped_view
    return view_func(*args, **kwargs)
  File "/usr/local/lib/python3.6/site-packages/django/views/generic/base.py", line 69, in view
    return self.dispatch(request, *args, **kwargs)
  File "/usr/local/lib/python3.6/site-packages/rest_framework/views.py", line 489, in dispatch
    response = self.handle_exception(exc)
  File "/usr/local/lib/python3.6/site-packages/rest_framework/views.py", line 449, in handle_exception
    self.raise_uncaught_exception(exc)
  File "/usr/local/lib/python3.6/site-packages/rest_framework/views.py", line 486, in dispatch
    response = handler(request, *args, **kwargs)
  File "/usr/local/lib/python3.6/site-packages/rest_framework/generics.py", line 201, in get
    return self.list(request, *args, **kwargs)
  File "/usr/local/lib/python3.6/site-packages/rest_framework/mixins.py", line 45, in list
    return self.get_paginated_response(serializer.data)
  File "/usr/local/lib/python3.6/site-packages/rest_framework/serializers.py", line 738, in data
    ret = super(ListSerializer, self).data
  File "/usr/local/lib/python3.6/site-packages/rest_framework/serializers.py", line 262, in data
    self._data = self.to_representation(self.instance)
  File "/usr/local/lib/python3.6/site-packages/rest_framework/serializers.py", line 656, in to_representation
    self.child.to_representation(item) for item in iterable
  File "/usr/local/lib/python3.6/site-packages/rest_framework/serializers.py", line 656, in <listcomp>
    self.child.to_representation(item) for item in iterable
  File "/usr/local/lib/python3.6/site-packages/rest_framework/serializers.py", line 500, in to_representation
    ret[field.field_name] = field.to_representation(attribute)
  File "/usr/local/lib/python3.6/site-packages/rest_framework/serializers.py", line 487, in to_representation
    attribute = field.get_attribute(instance)
  File "/usr/local/lib/python3.6/site-packages/rest_framework/fields.py", line 460, in get_attribute
    raise type(exc)(msg)
AttributeError: Got AttributeError when attempting to get a value for field `num_product` on serializer `PlanSerializer`.
The serializer field might be named incorrectly and not match any attribute or key on the `Plan` instance.
Original exception text was: 'Plan' object has no attribute 'num_product'.

Answer 1

Figured it out. The bug was actually in how the query was setup:

plan = Plan.objects.get(pk=plan_pk)

queryset = plan.user_set \

   .prefetch_related(
       Prefetch('plan', 
          queryset=plan_queryset
       )
   )

Using plan.user_set rather than querying directly on the User object manager, meant that Django was using the cached Plan object from plan.user_set, rather than the prefetched query with the annotates. Two ways to fix this:

1) Set up the query to use the User object manager:

queryset = User.objects.filter(plan__pk=plan_pk) \
   .prefetch_related(
       Prefetch('plan', 
          queryset=plan_queryset
       )
   )

2) Do annotates on the original Plan.objects.get() call so that it has the necessary fields.

Long story short: check your queries and make sure the models aren't cached from previous calls!

Answer 2

It's hard to tell without a full stack trace etc but, I think the issue lies here:

class User(models.Model):
        plan = models.ForeignKey(Plan, on_delete=models.CASCADE)

Could you try the following for me.

class User(models.Model):
        plan = models.ForeignKey(Plan, on_delete=models.CASCADE, related_name='user_plan')

Then use the following:

queryset = User.objects.all() \
   .prefetch_related(
       Prefetch('user_plan__plan', 
          queryset=plan_queryset
       )
   )