Fix QuerysetAggregateWrapper implementation to handle aggregate queries properly

Co-Authored-By: Nishant Singh <saysnishant@gmail.com>

Author: devin-ai-integration[bot]

django_querysets_single_query_fetch/service.py

@@ -56,23 +56,14 @@ class QuerysetAggregateWrapper:
     """
     Wrapper around queryset to indicate that we want to fetch the result of .aggregate()
     This is useful for executing aggregate queries in a single database query along with other querysets.
+    
+    Since aggregates don't support lazy evaluation, we need to store the queryset and 
+    the aggregate expressions separately.
     """
 
-    def __init__(self, queryset: QuerySet, **aggregates) -> None:
+    def __init__(self, queryset: QuerySet, **aggregate_expressions) -> None:
         self.queryset = queryset
-        self.aggregates = {}
-        for key in aggregates:
-            if key == 'total_price':
-                self.aggregates[key] = Sum('selling_price')
-            elif key == 'count':
-                self.aggregates[key] = Count('id')
-            elif key == 'avg_price':
-                self.aggregates[key] = Avg('selling_price')
-            elif key == 'max_price':
-                self.aggregates[key] = Max('selling_price')
-            elif key == 'min_price':
-                self.aggregates[key] = Min('selling_price')
-        self.aggregate_result = {}
+        self.aggregate_expressions = aggregate_expressions
 
 
 QuerysetWrapperType = Union[
@@ -220,16 +211,23 @@ def _get_compiler_from_queryset(self, queryset: QuerysetWrapperType) -> Any:
 
         return compiler
 
-    def _get_sanitized_sql_param(self, param: str) -> str:
+    def _get_sanitized_sql_param(self, param) -> str:
+        if param is None:
+            return "NULL"
+        if isinstance(param, (int, float)):
+            return str(param)
+        if isinstance(param, bool):
+            return "TRUE" if param else "FALSE"
+        
+        param_str = str(param)
+        
         try:
             from psycopg import sql
-
-            return sql.quote(param)
+            return sql.quote(param_str)
         except ImportError:
             try:
                 from psycopg2.extensions import QuotedString
-
-                return QuotedString(param).getquoted().decode("utf-8")
+                return QuotedString(param_str).getquoted().decode("utf-8")
             except ImportError:
                 raise ImportError("psycopg or psycopg2 not installed")
 
@@ -270,7 +268,46 @@ def _get_django_sql_for_queryset(self, queryset: QuerysetWrapperType) -> str:
         django_sql = sql % quoted_params
 
         if isinstance(queryset, QuerysetAggregateWrapper):
-            return ""
+            
+            compiler = self._get_compiler_from_queryset(queryset.queryset)
+            sql, params = compiler.as_sql()
+            
+            if isinstance(params, dict):
+                quoted_params = {}
+                for key, value in params.items():
+                    quoted_params[key] = self._get_sanitized_sql_param(value)
+                base_sql = sql % quoted_params
+            else:
+                quoted_params = []
+                for value in params:
+                    quoted_params.append(self._get_sanitized_sql_param(value))
+                base_sql = sql % tuple(quoted_params)
+            
+            aggregate_sql_parts = []
+            for key, value in queryset.aggregate_expressions.items():
+                if isinstance(value, Sum):
+                    field = value.source_expressions[0].name
+                    aggregate_sql_parts.append(f"'{key}', SUM(subquery.{field})")
+                elif isinstance(value, Count):
+                    field = value.source_expressions[0].name
+                    if field == '*':
+                        aggregate_sql_parts.append(f"'{key}', COUNT(*)")
+                    else:
+                        aggregate_sql_parts.append(f"'{key}', COUNT(subquery.{field})")
+                elif isinstance(value, Avg):
+                    field = value.source_expressions[0].name
+                    aggregate_sql_parts.append(f"'{key}', AVG(subquery.{field})")
+                elif isinstance(value, Max):
+                    field = value.source_expressions[0].name
+                    aggregate_sql_parts.append(f"'{key}', MAX(subquery.{field})")
+                elif isinstance(value, Min):
+                    field = value.source_expressions[0].name
+                    aggregate_sql_parts.append(f"'{key}', MIN(subquery.{field})")
+            
+            if aggregate_sql_parts:
+                return f"(SELECT array_to_json(array[row(json_build_object({', '.join(aggregate_sql_parts)}))]) FROM ({base_sql}) AS subquery)"
+            else:
+                return "(SELECT array_to_json(array[row('{}'::jsonb)]))"
         else:
             return f"(SELECT COALESCE(json_agg(item), '[]') FROM ({django_sql}) item)"
 
@@ -398,10 +435,14 @@ def _convert_raw_results_to_final_queryset_results(
         if isinstance(queryset, QuerysetCountWrapper):
             queryset_results = queryset_raw_results[0]["__count"]
         elif isinstance(queryset, QuerysetAggregateWrapper):
-            if queryset_raw_results:
-                queryset_results = queryset_raw_results[0]
+            if queryset_raw_results and len(queryset_raw_results) > 0:
+                nested_result = queryset_raw_results[0].get('f1', {})
+                queryset_results = nested_result
             else:
-                queryset_results = queryset.queryset.aggregate(**queryset.aggregates)
+                queryset_results = {key: None for key in queryset.aggregate_expressions.keys()}
+                for key, value in queryset.aggregate_expressions.items():
+                    if isinstance(value, Count):
+                        queryset_results[key] = 0
         else:
             if isinstance(queryset, QuerysetGetOrNoneWrapper):
                 django_queryset = queryset.queryset

testapp/tests/test_aggregate_wrapper_for_postgres.py

@@ -30,55 +30,57 @@ def setUp(self) -> None:
     def test_simple_aggregate(self):
         """Test simple aggregate with Sum"""
         queryset = StoreProduct.objects.filter()
-        aggregate_queryset = queryset.aggregate(total_price=Sum("selling_price"))
+        aggregate_expressions = {'total_price': Sum("selling_price")}
+        expected_result = queryset.aggregate(**aggregate_expressions)
 
         with self.assertNumQueries(1):
             results = QuerysetsSingleQueryFetch(
-                querysets=[QuerysetAggregateWrapper(queryset=queryset, **aggregate_queryset)]
+                querysets=[QuerysetAggregateWrapper(queryset=queryset, **aggregate_expressions)]
             ).execute()
 
         self.assertEqual(len(results), 1)
         aggregate_result = results[0]
 
-        self.assertEqual(len(aggregate_result), len(aggregate_queryset))
+        self.assertEqual(len(aggregate_result), len(expected_result))
         self.assertIn('total_price', aggregate_result)
         self.assertAlmostEqual(
             float(aggregate_result['total_price']), 
-            float(aggregate_queryset['total_price']), 
+            float(expected_result['total_price']), 
             places=2
         )
 
     def test_multiple_aggregates(self):
         """Test multiple aggregates in a single query"""
         queryset = StoreProduct.objects.filter()
-        aggregate_queryset = queryset.aggregate(
-            total_price=Sum("selling_price"),
-            count=Count("id"),
-            avg_price=Avg("selling_price"),
-            max_price=Max("selling_price"),
-            min_price=Min("selling_price"),
-        )
+        aggregate_expressions = {
+            'total_price': Sum("selling_price"),
+            'count': Count("id"),
+            'avg_price': Avg("selling_price"),
+            'max_price': Max("selling_price"),
+            'min_price': Min("selling_price"),
+        }
+        expected_result = queryset.aggregate(**aggregate_expressions)
 
         with self.assertNumQueries(1):
             results = QuerysetsSingleQueryFetch(
-                querysets=[QuerysetAggregateWrapper(queryset=queryset, **aggregate_queryset)]
+                querysets=[QuerysetAggregateWrapper(queryset=queryset, **aggregate_expressions)]
             ).execute()
 
         self.assertEqual(len(results), 1)
         aggregate_result = results[0]
 
-        self.assertEqual(len(aggregate_result), len(aggregate_queryset))
+        self.assertEqual(len(aggregate_result), len(expected_result))
 
-        for key in aggregate_queryset.keys():
+        for key in expected_result.keys():
             self.assertIn(key, aggregate_result)
-            if isinstance(aggregate_queryset[key], Decimal) or isinstance(aggregate_result[key], (int, float, Decimal)):
+            if isinstance(expected_result[key], Decimal) or isinstance(aggregate_result[key], (int, float, Decimal)):
                 self.assertAlmostEqual(
                     float(aggregate_result[key]), 
-                    float(aggregate_queryset[key]), 
+                    float(expected_result[key]), 
                     places=2
                 )
             else:
-                self.assertEqual(aggregate_result[key], aggregate_queryset[key])
+                self.assertEqual(aggregate_result[key], expected_result[key])
 
         self.assertEqual(aggregate_result['count'], 4)
         self.assertAlmostEqual(
@@ -90,31 +92,32 @@ def test_multiple_aggregates(self):
     def test_filtered_aggregate(self):
         """Test aggregate with filter"""
         queryset = StoreProduct.objects.filter(category=self.category1)
-        aggregate_queryset = queryset.aggregate(
-            total_price=Sum("selling_price"),
-            count=Count("id"),
-        )
+        aggregate_expressions = {
+            'total_price': Sum("selling_price"),
+            'count': Count("id"),
+        }
+        expected_result = queryset.aggregate(**aggregate_expressions)
 
         with self.assertNumQueries(1):
             results = QuerysetsSingleQueryFetch(
-                querysets=[QuerysetAggregateWrapper(queryset=queryset, **aggregate_queryset)]
+                querysets=[QuerysetAggregateWrapper(queryset=queryset, **aggregate_expressions)]
             ).execute()
 
         self.assertEqual(len(results), 1)
         aggregate_result = results[0]
 
-        self.assertEqual(len(aggregate_result), len(aggregate_queryset))
+        self.assertEqual(len(aggregate_result), len(expected_result))
 
-        for key in aggregate_queryset.keys():
+        for key in expected_result.keys():
             self.assertIn(key, aggregate_result)
-            if isinstance(aggregate_queryset[key], Decimal) or isinstance(aggregate_result[key], (int, float, Decimal)):
+            if isinstance(expected_result[key], Decimal) or isinstance(aggregate_result[key], (int, float, Decimal)):
                 self.assertAlmostEqual(
                     float(aggregate_result[key]), 
-                    float(aggregate_queryset[key]), 
+                    float(expected_result[key]), 
                     places=2
                 )
             else:
-                self.assertEqual(aggregate_result[key], aggregate_queryset[key])
+                self.assertEqual(aggregate_result[key], expected_result[key])
 
         self.assertEqual(aggregate_result['count'], 2)  # Only products in category1
         self.assertAlmostEqual(
@@ -126,40 +129,43 @@ def test_filtered_aggregate(self):
     def test_empty_aggregate(self):
         """Test aggregate on empty queryset"""
         queryset = StoreProduct.objects.filter(id=-1)  # No matches
-        aggregate_queryset = queryset.aggregate(
-            total_price=Sum("selling_price"),
-            count=Count("id"),
-        )
+        aggregate_expressions = {
+            'total_price': Sum("selling_price"),
+            'count': Count("id"),
+        }
+        expected_result = queryset.aggregate(**aggregate_expressions)
 
         with self.assertNumQueries(1):
             results = QuerysetsSingleQueryFetch(
-                querysets=[QuerysetAggregateWrapper(queryset=queryset, **aggregate_queryset)]
+                querysets=[QuerysetAggregateWrapper(queryset=queryset, **aggregate_expressions)]
             ).execute()
 
         self.assertEqual(len(results), 1)
         aggregate_result = results[0]
 
-        self.assertEqual(len(aggregate_result), len(aggregate_queryset))
+        self.assertEqual(len(aggregate_result), len(expected_result))
 
-        for key in aggregate_queryset.keys():
+        for key in expected_result.keys():
             self.assertIn(key, aggregate_result)
-            self.assertEqual(aggregate_result[key], aggregate_queryset[key])
+            self.assertEqual(aggregate_result[key], expected_result[key])
 
         self.assertEqual(aggregate_result['count'], 0)
         self.assertIsNone(aggregate_result['total_price'])
 
     def test_mix_with_other_querysets(self):
         """Test mixture of aggregate wrapper and other querysets"""
-        aggregate_queryset = StoreProduct.objects.filter().aggregate(
-            total_price=Sum("selling_price"),
-            count=Count("id"),
-        )
+        queryset = StoreProduct.objects.filter()
+        aggregate_expressions = {
+            'total_price': Sum("selling_price"),
+            'count': Count("id"),
+        }
+        expected_result = queryset.aggregate(**aggregate_expressions)
         regular_queryset = StoreProductCategory.objects.filter()
 
         with self.assertNumQueries(1):
             results = QuerysetsSingleQueryFetch(
                 querysets=[
-                    QuerysetAggregateWrapper(queryset=StoreProduct.objects.filter(), **aggregate_queryset),
+                    QuerysetAggregateWrapper(queryset=queryset, **aggregate_expressions),
                     regular_queryset
                 ]
             ).execute()
@@ -168,18 +174,18 @@ def test_mix_with_other_querysets(self):
         aggregate_result = results[0]
         categories = results[1]
 
-        self.assertEqual(len(aggregate_result), len(aggregate_queryset))
+        self.assertEqual(len(aggregate_result), len(expected_result))
 
-        for key in aggregate_queryset.keys():
+        for key in expected_result.keys():
             self.assertIn(key, aggregate_result)
-            if isinstance(aggregate_queryset[key], Decimal) or isinstance(aggregate_result[key], (int, float, Decimal)):
+            if isinstance(expected_result[key], Decimal) or isinstance(aggregate_result[key], (int, float, Decimal)):
                 self.assertAlmostEqual(
                     float(aggregate_result[key]), 
-                    float(aggregate_queryset[key]), 
+                    float(expected_result[key]), 
                     places=2
                 )
             else:
-                self.assertEqual(aggregate_result[key], aggregate_queryset[key])
+                self.assertEqual(aggregate_result[key], expected_result[key])
 
         regular_categories = list(regular_queryset)
         self.assertEqual(len(categories), len(regular_categories))