Search code examples
python-3.xdjangodjango-rest-frameworkannotationsmany-to-many

Annotate a Django Queryset using a 'related_name' across a ManyToMany relationship (NO 'Count', 'Avg', 'Max', ...)


I have models in my code similar to the following:

class CompanyProject(models.Model):
    """ This class holds project related information read in from the 'project'
        custom 'manage.py' command.
    """

    project_number = models.IntegerField(blank=False, null=False, unique=True)
    project_worktype = models.CharField(blank=True, max_length=255, null=True)
    created_at = models.DateTimeField(auto_now_add=True, blank=False, null=False)
    updated_at = models.DateTimeField(auto_now=True, blank=False, null=False)
    last_seen = models.DateField(blank=False, null=False)


    def get_project_subtypes(self):
        subtypes = self.project_subtype.all()
        return [ subtype.project_subtype for subtype in subtypes ]

    class Meta:
        ordering = ['project_number']

class CompanySubType(models.Model):
    class CompanySubTypeChoices(models.TextChoices):
        G1A = '1A', _('1A')
        G1B = '1B', _('1B')
        G2A = '2A', _('2A')
        G2B = '2B', _('2B')
        G3A = '3A', _('3A')
        G3B = '3B', _('3B')

    company_project = models.ManyToManyField(CompanyProject, related_name='project_subtype')
    project_subtype = models.CharField(blank=False, choices=CompanySubTypeChoices.choices, max_length=2, null=False)


class ListEntry(models.Model):

    list_project = models.OneToOneField(CompanyProject, on_delete=models.CASCADE, related_name='list_project')
    list_reviewer = models.ForeignKey('auth.User', on_delete=models.CASCADE, related_name='+')

I would like to return a set of annotated ListEntry objects annotated with a list of ALL of the project subtypes identified with the ListEntry project.

Eventually, I need to be able to pass this annotated data to a Django REST Framework serializer but I need to get the annotation working like I want it to first.

My problem is that I can annotate just fine doing something like this:

list_entry_qs = ListEntry.objects.prefetch_related(
    'list_project', 'list_reviewer'
).annotate(
    subtypes=F('list_pmatt__project_subtype__project_subtype')
).all()

and it works just fine. The issue I am having is that the query set that's returned from that command duplicates the list_entry object if it has more than one subtype annotation.

For example, if a project numbered 1234 has two project subtypes, '1A' and '3A', I get two list_entry objects: one with subtype annotation '1A' and a separate list_entry object for the same project with subtype annotation '3A'.

I would like a SINGLE object returned for project 1234 that has a subtype annotation of '1A', '2A' that will eventually be serialized to a JSON object.

There has to be an easy way to do this?


Solution

  • OK, After spending literally 8 hours on this, I finally found out how to do it:

    from django.contrib.postgres.aggregates.general import ArrayAgg
    
    q = ListEntry.objects.prefetch_related(
        'list_pmatt', 'list_reviewer', 'list_pmatt__project_subtype'
    ).annotate(
        my_subtype=ArrayAgg('list_pmatt__project_subtype__project_subtype')
    )
    
    >>> q[2].my_subtype
    ['1A', '3A']