Аннотирование нескольких внешних ключей одного типа в одной модели

У меня есть следующие модели. Мне абсолютно необходимо использовать несколько внешних ключей вместо поля "многие ко многим".

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)

Я хочу получить общую зарплату по каждому месту работы в следующем формате: {'name': 'MyCompany', 'admin_total': 5000, 'architect_total': 3000, 'engineer_total': 2000}. Я делаю это путем итерации по каждому из десяти сотрудников, проверяя их роли и складывая их вместе, если у них одинаковые роли:

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"),
)

Как видите, это очень длинный запрос, и он включает только одну из трех общих зарплат. А если добавить еще одну работу, то аннотация станет еще длиннее. Есть ли более эффективный способ сделать это?

Да, вы можете использовать функцию Django aggregate для получения общей зарплаты для каждой работы более эффективным способом. Вот как вы можете это сделать:

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}

Обратите внимание, что этот код предполагает, что поле job в модели Employee является полем выбора и использует строчную метку каждого job для создания ключей в конечном результате.

Вернуться на верх