Search code examples
djangodjango-annotate

How to annotate count filtered on datetime?


I've got Categories and Items. The items have an end field (datetime). Now I need to list all categories and display the related item count and the item count of items in the future. As an example:

  • Cat Foo, 2 items, 1 in the future.
  • Cat n, n items n in the future.

The list will be big. So the database has to do the heavy lifting and annotate both item_count and future_item_count.

Models:

from django.db import models

class Cat(models.Model):
    title = models.CharField(max_length=200)

class Item(models.Model):
    cat = models.ForeignKey(Cat)
    title = models.CharField(max_length=200)
    end = models.DateTimeField()

Create a category and two related items. One in the past, one in the future:

from datetime import timedelta
from django.utils import timezone

cat = Cat(title='Cat 1')
cat.save()

item_1 = Item(cat=cat, title="Item 1", end=timezone.now() - timedelta(days=1))
item_1.save()

item_2 = Item(cat=cat, title="Item 2", end=timezone.now() + timedelta(days=1))
item_2.save()

When I annotate item_count it works as expected:

from django.db.models import Count

Cat.objects.all().annotate(
    item_count=Count('item')).values('title', 'item_count')
# [{'item_count': 2, 'title': u'Cat 1'}]

I can NOT annotate filtered by Item.end (datetime). Is this possible with Django queries at all?

Cat.objects.all().annotate(
    item_count=Count('item'),
    future_item_count=Count('item').filter(
        end__gt=timezone.now())
    ).values(
        'title', 
        'item_count', 
        'future_item_count'
    )
# AttributeError: 'Count' object has no attribute 'filter'

I expect to get: [{'item_count': 2, 'future_item_count': 1, 'title': u'Cat 1'}]

I also tried RawSQL but lack SQL skills:

from django.db.models.expressions import RawSQL

Cat.objects.all().annotate(
    item_count=Count('item'), 
    future_item_count=RawSQL(
        """SELECT COUNT(*) 
           FROM project_item 
           JOIN project_item 
           AS foo 
           ON foo.cat_id = project_cat.id 
           WHERE project_item.end < NOW()""",
        ""
    )).values(
        'title', 
        'item_count', 
        'future_item_count'
    )
# [{'item_count': 2, 'future_item_count': 2L, 'title': u'Cat 1'}]

But when I change WHERE project_item.end < NOW()" in WHERE project_item.end > NOW()" I get the same result:

[{'item_count': 2, 'future_item_count': 2L, 'title': u'Cat 1'}]

How to format the raw SQL? Or can this be done with Django queries?


Solution

  • I personally haven't used RawSQL (still do things with .extra), but I think you don't need to JOIN project_item in your RawSQL statement. Just try with:

    RawSQL("""SELECT COUNT(*) 
                FROM project_item 
                WHERE 
                    project_item.cat_id = project_cat.id 
                    AND project_item.end < NOW()
    """)
    

    And one more thing I think you should not use .values AFTER .annotate, but BEFORE the annotation. So your full QuerySet should look like this:

    Cat.objects.values('title')\
        .annotate(
            item_count=Count('item'),
            future_item_count=RawSQL("""
                SELECT COUNT(*) 
                FROM project_item 
                WHERE 
                    project_item.cat_id = project_cat.id
                    AND project_item.end < NOW()
                """)
        )