Fixed #14030 -- Allowed annotations to accept all expressions

tests/aggregation/tests.py

@@ -3,12 +3,21 @@ from __future__ import unicode_literals
 import datetime
 from decimal import Decimal
 import re
+import warnings
 
+from django.core.exceptions import FieldError
 from django.db import connection
-from django.db.models import Avg, Sum, Count, Max, Min
+from django.db.models import (
+    Avg, Sum, Count, Max, Min,
+    Aggregate, F, Value, Func,
+    IntegerField, FloatField, DecimalField)
+with warnings.catch_warnings(record=True) as w:
+    warnings.simplefilter("always")
+    from django.db.models.sql import aggregates as sql_aggregates
 from django.test import TestCase
 from django.test.utils import Approximate
 from django.test.utils import CaptureQueriesContext
+from django.utils.deprecation import RemovedInDjango20Warning
 
 from .models import Author, Publisher, Book, Store
 
@@ -678,3 +687,271 @@ class BaseAggregateTestCase(TestCase):
             else:
                 self.assertNotIn('order by', qstr)
             self.assertEqual(qstr.count(' join '), 0)
+
+
+class ComplexAggregateTestCase(TestCase):
+    fixtures = ["aggregation.json"]
+
+    def test_nonaggregate_aggregation_throws(self):
+        with self.assertRaisesRegexp(TypeError, 'fail is not an aggregate expression'):
+            Book.objects.aggregate(fail=F('price'))
+
+    def test_nonfield_annotation(self):
+        book = Book.objects.annotate(val=Max(Value(2, output_field=IntegerField())))[0]
+        self.assertEqual(book.val, 2)
+        book = Book.objects.annotate(val=Max(Value(2), output_field=IntegerField()))[0]
+        self.assertEqual(book.val, 2)
+
+    def test_missing_output_field_raises_error(self):
+        with self.assertRaisesRegexp(FieldError, 'Cannot resolve expression type, unknown output_field'):
+            Book.objects.annotate(val=Max(Value(2)))[0]
+
+    def test_annotation_expressions(self):
+        authors = Author.objects.annotate(combined_ages=Sum(F('age') + F('friends__age'))).order_by('name')
+        authors2 = Author.objects.annotate(combined_ages=Sum('age') + Sum('friends__age')).order_by('name')
+        for qs in (authors, authors2):
+            self.assertEqual(len(qs), 9)
+            self.assertQuerysetEqual(
+                qs, [
+                    ('Adrian Holovaty', 132),
+                    ('Brad Dayley', None),
+                    ('Jacob Kaplan-Moss', 129),
+                    ('James Bennett', 63),
+                    ('Jeffrey Forcier', 128),
+                    ('Paul Bissex', 120),
+                    ('Peter Norvig', 103),
+                    ('Stuart Russell', 103),
+                    ('Wesley J. Chun', 176)
+                ],
+                lambda a: (a.name, a.combined_ages)
+            )
+
+    def test_aggregation_expressions(self):
+        a1 = Author.objects.aggregate(av_age=Sum('age') / Count('*'))
+        a2 = Author.objects.aggregate(av_age=Sum('age') / Count('age'))
+        a3 = Author.objects.aggregate(av_age=Avg('age'))
+        self.assertEqual(a1, {'av_age': 37})
+        self.assertEqual(a2, {'av_age': 37})
+        self.assertEqual(a3, {'av_age': Approximate(37.4, places=1)})
+
+    def test_order_of_precedence(self):
+        p1 = Book.objects.filter(rating=4).aggregate(avg_price=(Avg('price') + 2) * 3)
+        self.assertEqual(p1, {'avg_price': Approximate(148.18, places=2)})
+
+        p2 = Book.objects.filter(rating=4).aggregate(avg_price=Avg('price') + 2 * 3)
+        self.assertEqual(p2, {'avg_price': Approximate(53.39, places=2)})
+
+    def test_combine_different_types(self):
+        with self.assertRaisesRegexp(FieldError, 'Expression contains mixed types. You must set output_field'):
+            Book.objects.annotate(sums=Sum('rating') + Sum('pages') + Sum('price')).get(pk=4)
+
+        b1 = Book.objects.annotate(sums=Sum(F('rating') + F('pages') + F('price'),
+                                   output_field=IntegerField())).get(pk=4)
+        self.assertEqual(b1.sums, 383)
+
+        b2 = Book.objects.annotate(sums=Sum(F('rating') + F('pages') + F('price'),
+                                   output_field=FloatField())).get(pk=4)
+        self.assertEqual(b2.sums, 383.69)
+
+        b3 = Book.objects.annotate(sums=Sum(F('rating') + F('pages') + F('price'),
+                                   output_field=DecimalField(max_digits=6, decimal_places=2))).get(pk=4)
+        self.assertEqual(b3.sums, Decimal("383.69"))
+
+    def test_complex_aggregations_require_kwarg(self):
+        with self.assertRaisesRegexp(TypeError, 'Complex expressions require an alias'):
+            Author.objects.annotate(Sum(F('age') + F('friends__age')))
+        with self.assertRaisesRegexp(TypeError, 'Complex aggregates require an alias'):
+            Author.objects.aggregate(Sum('age') / Count('age'))
+
+    def test_aggregate_over_complex_annotation(self):
+        qs = Author.objects.annotate(
+            combined_ages=Sum(F('age') + F('friends__age')))
+
+        age = qs.aggregate(max_combined_age=Max('combined_ages'))
+        self.assertEqual(age['max_combined_age'], 176)
+
+        age = qs.aggregate(max_combined_age_doubled=Max('combined_ages') * 2)
+        self.assertEqual(age['max_combined_age_doubled'], 176 * 2)
+
+        age = qs.aggregate(
+            max_combined_age_doubled=Max('combined_ages') + Max('combined_ages'))
+        self.assertEqual(age['max_combined_age_doubled'], 176 * 2)
+
+        age = qs.aggregate(
+            max_combined_age_doubled=Max('combined_ages') + Max('combined_ages'),
+            sum_combined_age=Sum('combined_ages'))
+        self.assertEqual(age['max_combined_age_doubled'], 176 * 2)
+        self.assertEqual(age['sum_combined_age'], 954)
+
+        age = qs.aggregate(
+            max_combined_age_doubled=Max('combined_ages') + Max('combined_ages'),
+            sum_combined_age_doubled=Sum('combined_ages') + Sum('combined_ages'))
+        self.assertEqual(age['max_combined_age_doubled'], 176 * 2)
+        self.assertEqual(age['sum_combined_age_doubled'], 954 * 2)
+
+    def test_values_annotation_with_expression(self):
+        # ensure the F() is promoted to the group by clause
+        qs = Author.objects.values('name').annotate(another_age=Sum('age') + F('age'))
+        a = qs.get(pk=1)
+        self.assertEqual(a['another_age'], 68)
+
+        qs = qs.annotate(friend_count=Count('friends'))
+        a = qs.get(pk=1)
+        self.assertEqual(a['friend_count'], 2)
+
+        qs = qs.annotate(combined_age=Sum('age') + F('friends__age')).filter(pk=1).order_by('-combined_age')
+        self.assertEqual(
+            list(qs), [
+                {
+                    "name": 'Adrian Holovaty',
+                    "another_age": 68,
+                    "friend_count": 1,
+                    "combined_age": 69
+                },
+                {
+                    "name": 'Adrian Holovaty',
+                    "another_age": 68,
+                    "friend_count": 1,
+                    "combined_age": 63
+                }
+            ]
+        )
+
+        vals = qs.values('name', 'combined_age')
+        self.assertEqual(
+            list(vals), [
+                {
+                    "name": 'Adrian Holovaty',
+                    "combined_age": 69
+                },
+                {
+                    "name": 'Adrian Holovaty',
+                    "combined_age": 63
+                }
+            ]
+        )
+
+    def test_annotate_values_aggregate(self):
+        alias_age = Author.objects.annotate(
+            age_alias=F('age')
+        ).values(
+            'age_alias',
+        ).aggregate(sum_age=Sum('age_alias'))
+
+        age = Author.objects.values('age').aggregate(sum_age=Sum('age'))
+
+        self.assertEqual(alias_age['sum_age'], age['sum_age'])
+
+    def test_annotate_over_annotate(self):
+        author = Author.objects.annotate(
+            age_alias=F('age')
+        ).annotate(
+            sum_age=Sum('age_alias')
+        ).get(pk=1)
+
+        other_author = Author.objects.annotate(
+            sum_age=Sum('age')
+        ).get(pk=1)
+
+        self.assertEqual(author.sum_age, other_author.sum_age)
+
+    def test_annotated_aggregate_over_annotated_aggregate(self):
+        with self.assertRaisesRegexp(FieldError, "Cannot compute Sum\('id__max'\): 'id__max' is an aggregate"):
+            Book.objects.annotate(Max('id')).annotate(Sum('id__max'))
+
+    def test_add_implementation(self):
+        try:
+            # test completely changing how the output is rendered
+            def lower_case_function_override(self, qn, connection):
+                sql, params = qn.compile(self.source_expressions[0])
+                substitutions = dict(function=self.function.lower(), expressions=sql)
+                substitutions.update(self.extra)
+                return self.template % substitutions, params
+            setattr(Sum, 'as_' + connection.vendor, lower_case_function_override)
+
+            qs = Book.objects.annotate(sums=Sum(F('rating') + F('pages') + F('price'),
+                                       output_field=IntegerField()))
+            self.assertEqual(str(qs.query).count('sum('), 1)
+            b1 = qs.get(pk=4)
+            self.assertEqual(b1.sums, 383)
+
+            # test changing the dict and delegating
+            def lower_case_function_super(self, qn, connection):
+                self.extra['function'] = self.function.lower()
+                return super(Sum, self).as_sql(qn, connection)
+            setattr(Sum, 'as_' + connection.vendor, lower_case_function_super)
+
+            qs = Book.objects.annotate(sums=Sum(F('rating') + F('pages') + F('price'),
+                                       output_field=IntegerField()))
+            self.assertEqual(str(qs.query).count('sum('), 1)
+            b1 = qs.get(pk=4)
+            self.assertEqual(b1.sums, 383)
+
+            # test overriding all parts of the template
+            def be_evil(self, qn, connection):
+                substitutions = dict(function='MAX', expressions='2')
+                substitutions.update(self.extra)
+                return self.template % substitutions, ()
+            setattr(Sum, 'as_' + connection.vendor, be_evil)
+
+            qs = Book.objects.annotate(sums=Sum(F('rating') + F('pages') + F('price'),
+                                       output_field=IntegerField()))
+            self.assertEqual(str(qs.query).count('MAX('), 1)
+            b1 = qs.get(pk=4)
+            self.assertEqual(b1.sums, 2)
+        finally:
+            delattr(Sum, 'as_' + connection.vendor)
+
+    def test_complex_values_aggregation(self):
+        max_rating = Book.objects.values('rating').aggregate(
+            double_max_rating=Max('rating') + Max('rating'))
+        self.assertEqual(max_rating['double_max_rating'], 5 * 2)
+
+        max_books_per_rating = Book.objects.values('rating').annotate(
+            books_per_rating=Count('id') + 5
+        ).aggregate(Max('books_per_rating'))
+        self.assertEqual(
+            max_books_per_rating,
+            {'books_per_rating__max': 3 + 5})
+
+    def test_expression_on_aggregation(self):
+
+        # Create a plain expression
+        class Greatest(Func):
+            function = 'GREATEST'
+
+            def as_sqlite(self, qn, connection):
+                return super(Greatest, self).as_sql(qn, connection, function='MAX')
+
+        qs = Publisher.objects.annotate(
+            price_or_median=Greatest(Avg('book__rating'), Avg('book__price'))
+        ).filter(price_or_median__gte=F('num_awards')).order_by('pk')
+        self.assertQuerysetEqual(
+            qs, [1, 2, 3, 4], lambda v: v.pk)
+
+        qs2 = Publisher.objects.annotate(
+            rating_or_num_awards=Greatest(Avg('book__rating'), F('num_awards'),
+                                          output_field=FloatField())
+        ).filter(rating_or_num_awards__gt=F('num_awards')).order_by('pk')
+        self.assertQuerysetEqual(
+            qs2, [1, 2], lambda v: v.pk)
+
+    def test_backwards_compatibility(self):
+
+        class SqlNewSum(sql_aggregates.Aggregate):
+            sql_function = 'SUM'
+
+        class NewSum(Aggregate):
+            name = 'Sum'
+
+            def add_to_query(self, query, alias, col, source, is_summary):
+                klass = SqlNewSum
+                aggregate = klass(
+                    col, source=source, is_summary=is_summary, **self.extra)
+                query.annotations[alias] = aggregate
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", RemovedInDjango20Warning)
+            qs = Author.objects.values('name').annotate(another_age=NewSum('age') + F('age'))
+            a = qs.get(pk=1)
+            self.assertEqual(a['another_age'], 68)