Fixed #34013 -- Added QuerySet.order_by() support for annotation transforms.

Thanks Eugene Morozov and Ben Nace for the reports.

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)