Conditional Aggregates in Django

There is currently a ticket in the Django issue tracker (#11305) proposing to add conditional counts to Django's aggregation framework.  While implementing the ability to create a sum of donations in the new Donors module of Voter Universe, I found myself having to tackle this problem myself.

I wound up basing my solution on this code linked as a comment to ticket #11305 by user Alex.  The problem I ran into was when trying to use it against a column that is a foreign key.  Take the following example code using the CountIf class:

condition = {
    'donations__status': "'received'",
    'donations__recipient__isnull': True
}
queryset.annotate(
    num_donations = CountIf(
        'donations__amount', 
        condition=condition,
    )
)

The intention of the donations__recipient__isnull condition is to find all donations that don't have a recipient which in Voter Universe are considered donations to the campaign (donations with a recipient are used to track donations to other campaigns).  This is important because I was trying to get a count of all donations for a given donor given to the campaign.  I ran into two problems.

SQL Escaping Values

First, notice the double quote and single quote around received.  It appears the CountIf class isn't properly sql escaping the value, but this simple hack of single quotes contained inside double quotes did the trick.  A better long term solution is to properly escape the values in the CountIf class, but this worked for now.

Foreign Key Conditions

The second problem was bigger and the solution a bit more hackish.  When the query is built, this is the COUNT statement that CountIf produces:

COUNT(CASE (`donors_donation`.`status` = 'received'  
            AND `donors_recipient`.`id` IS NULL)

The problem here is that the table donors_recipient is not joined into the query.  The CountIf implementation is using Django's ORM layer to build a where query clause, but doesn't handle if the where query clause requires a new join to be added to the query.  Thus, while the donors_recipient table exists in the database, it's not a part of this query causing a query error on execution.  After fighting with this for a while, I came up with a modified CountIf statement that can hack around this problem (and by hack, I mean hack!):

from django.db import models

class SQLCountIf(models.sql.aggregates.Aggregate):
    is_ordinal = True
    sql_function = 'COUNT'
    sql_template= '%(function)s(CASE %(condition)s WHEN true THEN 1 ELSE NULL END)'

class CountIf(models.Aggregate):
    name = 'COUNT'

    def add_to_query(self, query, alias, col, source, is_summary):
        sql, params = query.model._default_manager.filter(**self.extra['condition']).query.where.as_sql()
        for search, replace in self.extra.get('sql_replace', {}).items():
            sql = sql.replace(search, replace)
        self.extra['condition'] = sql % tuple(params)
        aggregate = SQLCountIf(col, source=source, is_summary=is_summary, **self.extra)
        query.aggregates[alias] = aggregate

The new part here is the addition of a new optional keyword argument called sql_replace.  This argument contains a dictionary with the search string as the key and the replacement value as the value.  Using this, I can get conditional counts against id values of conditional fields by doing something like this:

condition = {
    'donations__status': "'received'",
    'donations__recipient__isnull': True
}
sql_replace = {'`donors_recipient`.`id`': '`donors_donation`.`recipient_id`'}
queryset.annotate(
    num_donations = CountIf(
        'donations__amount', 
        condition=condition,
        sql_replace=sql_replace
    )
)

As I said, it's not pretty and highly hackish, but it does get the job done.  This allows you to skip the join and go straight off of the foreign key's id.

I'm sure there's a more elegant way to do this.  I will spend more time once I have some to figure out a cleaner model, but for now, if you need to do conditional counts on foreign key fields, you have an interim solution.