How to prefetch or subquery a deeply nested object with condition in Django ORM

user_1216183 0
Viewed 1771

There are theses models and relations :

Hours --FK--> Task --FK--> Project <--FK-- Period

class Hour(models.Model):
  date = models.DateField(...)
  task = models.ForeignKey(Task, ...)

class Task(models.Model):
  project = models.ForeignKey(Project, ...)
  
class Project(models.Model):
  pass

class Period(models.Model):
  project = models.ForeignKey(Project,...)
  start = models.DateField(...)
  end = models.DateField(...)


Summary :
 Hour has one task
 Task has one project
 Period has one project
 Hour has a date
 Period has a start date and a end date

For a given date and a given project there is one or none period possible

I want to populate a period field in Hour objects the same way it would be done with prefetch_related (using queryset)

I want to have something like this :

hours = Hour.objects.prefetch_period().all()
hours.first().period # Period(...)

Using a custom queryset method like this :

class HourQuerySet(models.query.QuerySet):
  def prefetch_related(self):
    return ???

For the moment I've only succeed doing this using annotate and Subquery, but I only manage to retrieve the period_id and not the prefetched period :

def inject_period(self):
    period_qs = (
        Period.objects.filter(
            project__tasks=OuterRef("task"), start__lte=OuterRef("date"), end__gte=OuterRef("date")
        )
        .values("id")[:1]
    )
    return self.annotate(period_id=Subquery(period_qs))
django-orm django python
2 Answers

I think the only way to do this currently would be to define the filter for the queryset of a Prefetch object.

from django.db.models import Prefetch
Hour.objects.prefetch_related(
    Prefetch(
        'task__project__periods',
        queryset=Period.objects.filter(
            project__tasks__date__gte=F('start'),
            project__tasks__date__lte=F('end'),
        ).distinct()
    )
)

This should effectively change your second query from doing something like

SELECT ... FROM app_period JOIN ... 
WHERE apps_period.project_id in (...)

to

SELECT DISTINCT ... FROM app_period JOIN ... 
WHERE apps_period.project_id in (...)
  AND apps_tasks.date BETWEEN apps_period.end AND apps_period.start
user_16373510

I have found the following solution, but it a bit hacky and I'm not sure to end using it.

I override the django's internal _fetch_all method of the QuerySet which is called when the queryset is fired. Then I do a custom prefetching and set attribute of the instances.

This probably need some futher optimisations.

class HourQuerySet(models.query.QuerySet):
  # With annotate and Subquery, search and define period_id
  def prefetch_period(self):
    period_qs = (
        Period.objects.filter(
            project__tasks=OuterRef("task"), start__lte=OuterRef("date"), end__gte=OuterRef("date")
        )
        .values("id")[:1]
    )
    return self.annotate(period_id=Subquery(period_qs))

  # Override _fetch_all method to manually prefetch and inject period in returned instances (for which who have a period_id defined)
  def _fetch_all(self):
    super()._fetch_all()
    if not self._result_cache or type(self._result_cache[0]) is dict:
        return
    period_ids = [r.period_id for r in self._result_cache]
    if not period_ids:
        return
    periods = {p.id: p for p in Period.objects.filter(id__in=period_ids)}
    for wh in self._result_cache:
        setattr(wh, "period", periods.get(wh.period_id))
user_12161830
Related