Annotate over multiple Foreign keys of the same type in one model

I have the following models. I absolutely have to use multiple foreign keys instead of a many-to-many field.

class Job(models.IntegerChoices):
    ADMIN = (0, "Admin")
    ARCHITECT = (1, "Architect")
    ENGINEER = (2, "Engineer")


class Employee(models.Model):
    job = models.IntegerField(_("Job"), choices=Job.choices)
    salary = models.DecimalField(_("Salary"), max_digits=12, decimal_places=4)


class Company(models.Model):
    name = models.CharField(...)
    employee_one = models.ForeignKey(Employee, on_delete=models.SET_NULL, null=True)
    employee_two = models.ForeignKey(Employee, on_delete=models.SET_NULL, null=True)
    employee_three = models.ForeignKey(Employee, on_delete=models.SET_NULL, null=True)
    ...
    employee_ten = models.ForeignKey(Employee, on_delete=models.SET_NULL, null=True)

I want to get the total salary for each job, as in the following format: {'name': 'MyCompany', 'admin_total': 5000, 'architect_total': 3000, 'engineer_total': 2000}. I do this by iterating through each of the ten employees, checking their role and adding them together if they have the same role:

Company.objects.all().annotate(
    admin_one=Case(
        When(employee_one__job=Job.ADMIN, then=F("employee_one__salary")),
        default=0,
        output_field=models.DecimalField(max_digits=12, decimal_places=4),
    ),
    admin_two=Case(
        When(employee_two__job=Job.ADMIN, then=F("employee_two__salary")),
        default=0,
        output_field=models.DecimalField(max_digits=12, decimal_places=4),
    ),
    ...,
    admin_total=F("admin_one") + F("admin_two") + ... + F("admin_ten"),
)

As you can see this is just a very long query and it is only including the one of the three total salaries. And if another job is added, the annotation will just get longer. Is there a more efficient way to do this?

Yes, you can use Django's aggregate function to get the total salary for each job in a more efficient way. Here's how you can do it:

from django.db.models import Sum

jobs_salary = Company.objects.filter(
    employee_one__job=Job.ADMIN
).aggregate(
    admin_total=Sum('employee_one__salary')
)

# Get total salary for each job:
total_salaries = {}
for job in Job:
    total = Company.objects.filter(
        Q(employee_one__job=job) | Q(employee_two__job=job) | Q(employee_three__job=job) |
        Q(employee_four__job=job) | Q(employee_five__job=job) | Q(employee_six__job=job) |
        Q(employee_seven__job=job) | Q(employee_eight__job=job) | Q(employee_nine__job=job) |
        Q(employee_ten__job=job)
    ).aggregate(
        total=Sum('employee_one__salary') + Sum('employee_two__salary') + Sum('employee_three__salary') +
        Sum('employee_four__salary') + Sum('employee_five__salary') + Sum('employee_six__salary') +
        Sum('employee_seven__salary') + Sum('employee_eight__salary') + Sum('employee_nine__salary') +
        Sum('employee_ten__salary')
    )
    total_salaries[job.label.lower() + '_total'] = total['total'] or 0

# Create the desired result:
result = {'name': 'MyCompany', **total_salaries}

Note that this code assumes that the job field in Employee model is a choice field and uses the lowercase label of each job to create the keys in the final result.

Back to Top