From b0ad41198b3e333f57351e3fce5a1fb47f23f376 Mon Sep 17 00:00:00 2001 From: Simon Charette Date: Fri, 8 Dec 2023 02:03:14 -0500 Subject: [PATCH] Fixed #34013 -- Added QuerySet.order_by() support for annotation transforms. Thanks Eugene Morozov and Ben Nace for the reports. --- django/db/models/sql/compiler.py | 26 ++++++++++------- docs/releases/5.1.txt | 3 ++ tests/aggregation/tests.py | 23 +++++++++++++++ .../comparison/test_json_object.py | 19 ++++++++++++- tests/postgres_tests/test_array.py | 10 +++++++ tests/queries/test_qs_combinators.py | 28 ++++++++++++++++++- 6 files changed, 97 insertions(+), 12 deletions(-) diff --git a/django/db/models/sql/compiler.py b/django/db/models/sql/compiler.py index cb559308cd..10338259d5 100644 --- a/django/db/models/sql/compiler.py +++ b/django/db/models/sql/compiler.py @@ -387,18 +387,24 @@ class SQLCompiler: True, ) continue - if col in self.query.annotations: - # References to an expression which is masked out of the SELECT - # clause. + + ref, *transforms = col.split(LOOKUP_SEP) + if expr := self.query.annotations.get(ref): if self.query.combinator and self.select: + if transforms: + raise NotImplementedError( + "Ordering combined queries by transforms is not " + "implemented." + ) # Don't use the resolved annotation because other - # combinated queries might define it differently. - expr = F(col) - else: - expr = self.query.annotations[col] - if isinstance(expr, Value): - # output_field must be resolved for constants. - expr = Cast(expr, expr.output_field) + # combined queries might define it differently. + expr = F(ref) + if transforms: + for name in transforms: + expr = self.query.try_transform(expr, name) + if isinstance(expr, Value): + # output_field must be resolved for constants. + expr = Cast(expr, expr.output_field) yield OrderBy(expr, descending=descending), False continue diff --git a/docs/releases/5.1.txt b/docs/releases/5.1.txt index 395acb5a89..1009ffe739 100644 --- a/docs/releases/5.1.txt +++ b/docs/releases/5.1.txt @@ -181,6 +181,9 @@ Models :class:`~django.db.models.expressions.ValueRange` allows excluding rows, groups, and ties from the window frames. +* :meth:`.QuerySet.order_by` now supports ordering by annotation transforms + such as ``JSONObject`` keys and ``ArrayAgg`` indices. + Requests and Responses ~~~~~~~~~~~~~~~~~~~~~~ diff --git a/tests/aggregation/tests.py b/tests/aggregation/tests.py index b01df88109..bedc2730a2 100644 --- a/tests/aggregation/tests.py +++ b/tests/aggregation/tests.py @@ -25,6 +25,7 @@ from django.db.models import ( Subquery, Sum, TimeField, + Transform, Value, Variance, When, @@ -1727,6 +1728,28 @@ class AggregateTestCase(TestCase): ordered=False, ) + def test_order_by_aggregate_transform(self): + class Mod100(Mod, Transform): + def __init__(self, expr): + super().__init__(expr, 100) + + sum_field = IntegerField() + sum_field.register_instance_lookup(Mod100, "mod100") + publisher_pages = ( + Book.objects.values("publisher") + .annotate(sum_pages=Sum("pages", output_field=sum_field)) + .order_by("sum_pages__mod100") + ) + self.assertQuerySetEqual( + publisher_pages, + [ + {"publisher": self.p2.id, "sum_pages": 528}, + {"publisher": self.p4.id, "sum_pages": 946}, + {"publisher": self.p1.id, "sum_pages": 747}, + {"publisher": self.p3.id, "sum_pages": 1482}, + ], + ) + def test_empty_result_optimization(self): with self.assertNumQueries(0): self.assertEqual( diff --git a/tests/db_functions/comparison/test_json_object.py b/tests/db_functions/comparison/test_json_object.py index 7a10657317..9a3d48288c 100644 --- a/tests/db_functions/comparison/test_json_object.py +++ b/tests/db_functions/comparison/test_json_object.py @@ -12,7 +12,12 @@ from ..models import Article, Author class JSONObjectTests(TestCase): @classmethod def setUpTestData(cls): - Author.objects.create(name="Ivan Ivanov", alias="iivanov") + Author.objects.bulk_create( + [ + Author(name="Ivan Ivanov", alias="iivanov"), + Author(name="Bertha Berthy", alias="bberthy"), + ] + ) def test_empty(self): obj = Author.objects.annotate(json_object=JSONObject()).first() @@ -88,6 +93,18 @@ class JSONObjectTests(TestCase): obj = Article.objects.annotate(json_object=JSONObject(text=F("text"))).first() self.assertEqual(obj.json_object, {"text": "x" * 4000}) + def test_order_by_key(self): + qs = Author.objects.annotate(attrs=JSONObject(alias=F("alias"))).order_by( + "attrs__alias" + ) + self.assertQuerySetEqual(qs, Author.objects.order_by("alias")) + + def test_order_by_nested_key(self): + qs = Author.objects.annotate( + attrs=JSONObject(nested=JSONObject(alias=F("alias"))) + ).order_by("-attrs__nested__alias") + self.assertQuerySetEqual(qs, Author.objects.order_by("-alias")) + @skipIfDBFeature("has_json_object_function") class JSONObjectNotSupportedTests(TestCase): diff --git a/tests/postgres_tests/test_array.py b/tests/postgres_tests/test_array.py index 48b03c626a..8aaa7be077 100644 --- a/tests/postgres_tests/test_array.py +++ b/tests/postgres_tests/test_array.py @@ -469,6 +469,16 @@ class TestQuerying(PostgreSQLTestCase): self.assertIn("GROUP BY 2", sql) self.assertIn("ORDER BY 2", sql) + def test_order_by_arrayagg_index(self): + qs = ( + NullableIntegerArrayModel.objects.values("order") + .annotate(ids=ArrayAgg("id")) + .order_by("-ids__0") + ) + self.assertQuerySetEqual( + qs, [{"order": obj.order, "ids": [obj.id]} for obj in reversed(self.objs)] + ) + def test_index(self): self.assertSequenceEqual( NullableIntegerArrayModel.objects.filter(field__0=2), self.objs[1:3] diff --git a/tests/queries/test_qs_combinators.py b/tests/queries/test_qs_combinators.py index 71d2418f2b..cb5ba1a269 100644 --- a/tests/queries/test_qs_combinators.py +++ b/tests/queries/test_qs_combinators.py @@ -1,7 +1,16 @@ import operator from django.db import DatabaseError, NotSupportedError, connection -from django.db.models import Exists, F, IntegerField, OuterRef, Subquery, Value +from django.db.models import ( + Exists, + F, + IntegerField, + OuterRef, + Subquery, + Transform, + Value, +) +from django.db.models.functions import Mod from django.test import TestCase, skipIfDBFeature, skipUnlessDBFeature from django.test.utils import CaptureQueriesContext @@ -322,6 +331,23 @@ class QuerySetSetOperationTests(TestCase): operator.itemgetter("num"), ) + def test_order_by_annotation_transform(self): + class Mod2(Mod, Transform): + def __init__(self, expr): + super().__init__(expr, 2) + + output_field = IntegerField() + output_field.register_instance_lookup(Mod2, "mod2") + qs1 = Number.objects.annotate( + annotation=Value(1, output_field=output_field), + ) + qs2 = Number.objects.annotate( + annotation=Value(2, output_field=output_field), + ) + msg = "Ordering combined queries by transforms is not implemented." + with self.assertRaisesMessage(NotImplementedError, msg): + list(qs1.union(qs2).order_by("annotation__mod2")) + def test_union_with_select_related_and_order(self): e1 = ExtraInfo.objects.create(value=7, info="e1") a1 = Author.objects.create(name="a1", num=1, extra=e1)