Django Aggregate Subquery

django_aggregate_subquery.py

"""
This can be used to get the aggregated value of a field in a subquery, like:

  >>> comments_subquery = Comment.objects.filter(
  ...    post=OuterRef("pk"),
  ... )
  >>> posts = Post.objects.annotate(
  ...     most_recent_comment_date=MinSubquery(
  ...         comments_subquery.values("created_at"),
  ...         "created_at",
  ...     )
  ... )
  
This would produce the following SQL:

  SELECT
    post.id,
    (SELECT MIN(updated_at) FROM comment WHERE comment.post_id = post.id) as most_recent_comment_date
  FROM post
  
"""

from typing import Any, Self, cast

from django.db import connection
from django.db.models.expressions import Subquery


class AggregateSubquery(Subquery):
    agg_op: str
    params: list[Any]

    def __init__(self, queryset, field: str, *args, **kwargs):
        self.agg_field = field
        self.params = []
        super().__init__(queryset, *args, **kwargs)

    def copy(self):
        clone = cast(Self, super().copy())
        clone.agg_field = self.agg_field
        clone.params = self.params
        return clone

    def as_sql(self, *args, **kwargs):
        sql, sql_params = super().as_sql(*args, **kwargs)
        variables = [f'__t."{self.agg_field}"', *self.params]
        sql = f'(SELECT {self.agg_op}({", ".join(variables)}) FROM ({sql}) AS __t)'
        return sql, sql_params


class CountSubquery(AggregateSubquery):
    agg_op = "COUNT"


class SumSubquery(AggregateSubquery):
    agg_op = "SUM"


class MinSubquery(AggregateSubquery):
    agg_op = "MIN"


class MaxSubquery(AggregateSubquery):
    agg_op = "MAX"


class AvgSubquery(AggregateSubquery):
    agg_op = "AVG"