Annotating a query with a values_list, using foreign key and many-to-many in Django

syberen 0
Viewed 29

I have a django operation that I'm performing in 2 steps, but I'd like to see if there's a way to do it with just 1 database query. Say I have these models:

class Budget(models.Model):
    ...

class Company(models.Model):
    budgets = models.ManyToManyField(
        Budget, related_name="companies"
    )

class Employee(models.Model):
    company = models.ForeignKey(
        Company, on_delete=models.CASCADE, related_name="employees"
    )

Now I want to get an employee, and find the IDs of the budgets that are associated with this employee, through the company:

employee = Employee.objects.get(id=employee_id)

allowed_budgets = employee.company.budgets.values_list("id", flat=True)

Am I correct that this would perform 2 database queries? Is there a way to add allowed_budgets as an annotation on Employee so only 1 query is performed?

employee = Employee.objects
     .annotate(
        allowed_budgets=# how do I get employee.company.budgets.values_list("id", flat=True) here?
     )
     .get(id=employee_id)

Thanks a lot!

django python
1 Answers

Maybe you can create a subclass of Subquery, that changes the SQL it outputs. Subqueries and annotations are powerful tools, which are made to be nested and composed into more complicated computations. For instance:

from django.db.models import OuterRef, Subquery,Count
Employee.objects.annotate(
    count=Subquery(
        Budget.objects
            .filter(budgets=OuterRef('id'))            
            .values('budgets')            
            .annotate(count=Count('id'))            
            .values('count')
    )
)

For PostgreSQL:

from django.contrib.postgres.aggregates import ArrayAgg, StringAgg
Employee.objects.annotate(
    cls=Subquery(
        Budgets.objects.filter(id=OuterRef('id')).values('id')\
        .annotate(budgets=ArrayAgg('budgets')).values('budgets')
    )
)

More information: https://django.readthedocs.io/en/stable/ref/models/expressions.html

Derya Cortuk0
Related