Fix annotation in ORDER BY of DISTINCT query (#1886)

Author: henadzit

CHANGELOG.rst

@@ -18,6 +18,7 @@ Added
 Fixed
 ^^^^^
 - Fix update pk field raises unfriendly error (#1873)
+- Using `.distinct()` with an annotation and `.order_by()` produces invalid SQL for PostgreSQL (#1886)
 
 Changed
 ^^^^^^^

tests/test_order_by.py

@@ -9,6 +9,7 @@
 from tortoise.contrib import test
 from tortoise.contrib.test.condition import NotEQ
 from tortoise.exceptions import ConfigurationError, FieldError
+from tortoise.expressions import Q, Case, When
 from tortoise.functions import Count, Sum
 
 
@@ -94,6 +95,45 @@ async def test_order_by_aggregation_reversed(self):
         )
         self.assertEqual([t.name for t in tournaments], ["1", "2"])
 
+    async def test_distinct_values_with_annotation(self):
+        await Tournament.create(name="3")
+        await Tournament.create(name="1")
+        await Tournament.create(name="2")
+
+        tournaments = (
+            await Tournament.annotate(
+                name_orderable=Case(
+                    When(Q(name="1"), then="1"),
+                    When(Q(name="2"), then="2"),
+                    When(Q(name="3"), then="3"),
+                    default="-1",
+                ),
+            )
+            .distinct()
+            .order_by("name_orderable", "-created")
+            .values("name", "name_orderable", "created")
+        )
+        self.assertEqual([t["name"] for t in tournaments], ["1", "2", "3"])
+
+    async def test_distinct_all_withl_annotation(self):
+        await Tournament.create(name="3")
+        await Tournament.create(name="1")
+        await Tournament.create(name="2")
+
+        tournaments = (
+            await Tournament.annotate(
+                name_orderable=Case(
+                    When(Q(name="1"), then="1"),
+                    When(Q(name="2"), then="2"),
+                    When(Q(name="3"), then="3"),
+                    default="-1",
+                ),
+            )
+            .distinct()
+            .order_by("name_orderable", "-created")
+        )
+        self.assertEqual([t.name for t in tournaments], ["1", "2", "3"])
+
 
 class TestDefaultOrdering(test.TestCase):
     @test.requireCapability(dialect=NotEQ("oracle"))

tests/test_values.py

@@ -4,7 +4,7 @@
 from tortoise.contrib import test
 from tortoise.contrib.test.condition import In, NotEQ
 from tortoise.exceptions import FieldError
-from tortoise.expressions import Function
+from tortoise.expressions import Q, Case, Function, When
 from tortoise.functions import Length, Trim
 
 
@@ -214,3 +214,41 @@ class TruncMonth(Function):
             sql,
             'SELECT DATE_FORMAT("created",?) "date" FROM "tournament"',
         )
+
+    async def test_order_by_annotation_not_in_values(self):
+        await Tournament.create(name="2")
+        await Tournament.create(name="3")
+        await Tournament.create(name="1")
+
+        tournaments = (
+            await Tournament.annotate(
+                name_orderable=Case(
+                    When(Q(name="1"), then="a"),
+                    When(Q(name="2"), then="b"),
+                    When(Q(name="3"), then="c"),
+                    default="z",
+                )
+            )
+            .order_by("name_orderable")
+            .values("name")
+        )
+        self.assertEqual([t["name"] for t in tournaments], ["1", "2", "3"])
+
+    async def test_order_by_annotation_not_in_values_list(self):
+        await Tournament.create(name="2")
+        await Tournament.create(name="3")
+        await Tournament.create(name="1")
+
+        tournaments = (
+            await Tournament.annotate(
+                name_orderable=Case(
+                    When(Q(name="1"), then="a"),
+                    When(Q(name="2"), then="b"),
+                    When(Q(name="3"), then="c"),
+                    default="z",
+                )
+            )
+            .order_by("name_orderable")
+            .values_list("name")
+        )
+        self.assertEqual(tournaments, [("1",), ("2",), ("3",)])

tortoise/queryset.py

@@ -1,15 +1,15 @@
 from __future__ import annotations
 
 import types
-from collections.abc import AsyncIterator, Callable, Generator, Iterable
+from collections.abc import AsyncIterator, Callable, Collection, Generator, Iterable
 from copy import copy
 from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar, cast, overload
 
 from pypika_tortoise import JoinType, Order, Table
 from pypika_tortoise.analytics import Count
 from pypika_tortoise.functions import Cast
 from pypika_tortoise.queries import QueryBuilder
-from pypika_tortoise.terms import Case, Field, Star, Term, ValueWrapper
+from pypika_tortoise.terms import Case, Field, Star, Term, ValueWrapper, PseudoColumn
 from typing_extensions import Literal, Protocol
 
 from tortoise.backends.base.client import BaseDBAsyncClient, Capabilities
@@ -179,7 +179,8 @@ def resolve_ordering(
         model: type[Model],
         table: Table,
         orderings: Iterable[tuple[str, str | Order]],
-        annotations: dict[str, Any],
+        annotations: dict[str, Term | Expression],
+        fields_for_select: Collection[str] | None = None,
     ) -> None:
         """
         Applies standard ordering to QuerySet.
@@ -189,6 +190,8 @@ def resolve_ordering(
             (to allow self referential joins)
         :param orderings: What columns/order to order by
         :param annotations:  Annotations that may be ordered on
+        :param fields_for_select: Contains fields that are selected in the SELECT clause if
+            .only(), .values() or .values_list() are used.
 
         :raises FieldError: If a field provided does not exist in model.
         """
@@ -214,18 +217,27 @@ def resolve_ordering(
                     {},
                 )
             elif field_name in annotations:
-                if isinstance(annotation := annotations[field_name], Term):
-                    term: Term = annotation
+                term: Term
+                if not fields_for_select or field_name in fields_for_select:
+                    # The annotation is SELECTed, we can just reference it in the following cases:
+                    # - Empty fields_for_select means that all columns and annotations are selected,
+                    #   hence we can reference the annotation.
+                    # - The annotation is in fields_for_select, hence we can reference it.
+                    term = PseudoColumn(field_name)
                 else:
-                    annotation_info = annotation.resolve(
-                        ResolveContext(
-                            model=self.model,
-                            table=table,
-                            annotations=annotations,
-                            custom_filters={},
-                        )
-                    )
-                    term = annotation_info.term
+                    # The annotation is not in SELECT, resolve it
+                    annotation = annotations[field_name]
+                    if isinstance(annotation, Term):
+                        term = annotation
+                    else:
+                        term = annotation.resolve(
+                            ResolveContext(
+                                model=self.model,
+                                table=table,
+                                annotations=annotations,
+                                custom_filters={},
+                            )
+                        ).term
                 self.query = self.query.orderby(term, order=ordering[1])
             else:
                 field_object = model._meta.fields_map.get(field_name)
@@ -1078,7 +1090,11 @@ def _make_query(self) -> None:
             if append_item not in self._select_related_idx:
                 self._select_related_idx.append(append_item)
         self.resolve_ordering(
-            self.model, self.model._meta.basetable, self._orderings, self._annotations
+            self.model,
+            self.model._meta.basetable,
+            self._orderings,
+            self._annotations,
+            self._fields_for_select,
         )
         self.resolve_filters()
         if self._limit is not None:
@@ -1562,6 +1578,7 @@ def _make_query(self) -> None:
             table=self.model._meta.basetable,
             orderings=self._orderings,
             annotations=self._annotations,
+            fields_for_select=self._fields_for_select_list,
         )
         self.resolve_filters()
         if self._limit:
@@ -1683,6 +1700,7 @@ def _make_query(self) -> None:
             table=self.model._meta.basetable,
             orderings=self._orderings,
             annotations=self._annotations,
+            fields_for_select=self._fields_for_select.keys(),
         )
         self.resolve_filters()