Fix DISTINCT ON for PostgreSQL: resolve source_field/relations, include Meta.ordering in validation, defer backend check.

Author: VladislavYar

tests/test_distinct.py

@@ -1,9 +1,10 @@
 import pytest
 
-from tests.testmodels import Tournament
+from tests.testmodels import Author, Book, DefaultOrdered, SourceFieldPk, Tournament
 from tortoise.contrib import test
 from tortoise.contrib.test.condition import NotIn
 from tortoise.exceptions import OperationalError
+from tortoise.functions import Count
 
 # ---------------------------------------------------------------------------
 # Basic DISTINCT (all databases)
@@ -185,6 +186,55 @@ async def test_distinct_on_only_with_order_by(db):
     assert [(t.name, t.desc) for t in tournaments] == [("1", "b"), ("2", "c")]
 
 
+@test.requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_distinct_on_filter_by_model(db):
+    tournament_1 = await Tournament.create(name="1", desc="a")
+    await Tournament.create(name="1", desc="b")
+    tournament_2 = await Tournament.create(name="2", desc="c")
+    tournaments = await Tournament.filter(name__in=["1", "2"]).distinct("name")
+    assert [tournament_1, tournament_2] == tournaments
+
+
+@test.requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_distinct_on_source_field(db):
+    await SourceFieldPk.create(name="1")
+    await SourceFieldPk.create(name="2")
+    await SourceFieldPk.all().distinct("id")
+
+
+@test.requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_distinct_on_annotate_by_model(db):
+    await Tournament.create(name="1", desc="a")
+    await Tournament.create(name="1", desc="b")
+    await Tournament.create(name="2", desc="c")
+    tournaments = (
+        await Tournament.annotate(count_name=Count("name")).distinct("name").order_by("name")
+    )
+    assert [1, 1] == [tournaments[0].count_name, tournaments[0].count_name]
+
+
+@test.requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_distinct_on_default_ordered(db):
+    await DefaultOrdered.create(one="1", second=1)
+    await DefaultOrdered.create(one="2", second=2)
+    await DefaultOrdered.all().distinct("one")
+
+
+@test.requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_distinct_on_by_relation(db):
+    author_1 = await Author.create(name="1")
+    author_2 = await Author.create(name="1")
+    await Book.create(name="1", rating=1, subject="1", author=author_1)
+    await Book.create(name="2", rating=2, subject="2", author=author_2)
+    books = await Book.all().distinct("author__name")
+    assert len(books) == 1
+
+
 # ---------------------------------------------------------------------------
 # DISTINCT ON validation errors
 # ---------------------------------------------------------------------------
@@ -202,4 +252,13 @@ async def test_distinct_on_invalid_order_by(db):
 @pytest.mark.asyncio
 async def test_distinct_on_not_supported_outside_postgres(db):
     with pytest.raises(OperationalError):
-        Tournament.all().distinct("name")
+        await Tournament.all().distinct("name")
+
+
+@test.requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_distinct_on_invalid_default_ordered(db):
+    await DefaultOrdered.create(one="1", second=1)
+    await DefaultOrdered.create(one="2", second=2)
+    with pytest.raises(OperationalError):
+        await DefaultOrdered.all().distinct("second")

tortoise/queryset.py

@@ -102,7 +102,7 @@ def _choose_db(self, for_write: bool = False) -> BaseDBAsyncClient:
             db = router.db_for_read(self.model)
         return db or self.model._meta.db
 
-    def _apply_db(self, db: BaseDBAsyncClient) -> None:
+    def _apply_db(self, db: BaseDBAsyncClient | None) -> None:
         """
         Set the database connection for this query and update the query builder dialect.
 
@@ -291,6 +291,63 @@ def resolve_ordering(
 
                 self.query = self.query.orderby(field, order=ordering[1])
 
+    def resolve_distinct(
+        self,
+        distinct: bool,
+        distinct_on: list[str],
+        orderings: Iterable[tuple[str, str | Order]],
+        annotations: dict[str, Term | Expression],
+    ) -> None:
+        if not distinct:
+            return
+        if not orderings and self.model._meta.ordering and not annotations:
+            orderings = self.model._meta.ordering
+        self.query._distinct = True
+        if distinct_on:
+            if not isinstance(self.query, PostgreSQLQueryBuilder):
+                raise OperationalError("DISTINCT ON is only supported by PostgreSQL")
+            ordering_fields = [ordering[0] for ordering in orderings]
+            len_ordering_fields = len(ordering_fields)
+            for i, field in enumerate(distinct_on):
+                if ordering_fields and (i >= len_ordering_fields or ordering_fields[i] != field):
+                    raise OperationalError(
+                        f"DISTINCT ON fields must match the leading ORDER BY fields. "
+                        f"Expected ORDER BY to start with {distinct_on!r}."
+                    )
+            self.query._distinct_on = []
+            distinct_on_by_source_field = []
+            for field_name in distinct_on:
+                field_object = self.model._meta.fields_map.get(field_name)
+                part_after = field_name
+                related_table = self.model._meta.basetable
+                related_model: type[Model] = self.model
+                while part_after:
+                    related_field_name, __, part_after = part_after.partition("__")
+                    if related_field_name in related_model._meta.fetch_fields:
+                        related_field = cast(
+                            RelationalField, self.model._meta.fields_map[related_field_name]
+                        )
+                        related_table = self._join_table_by_field(
+                            related_table, related_field_name, related_field
+                        )
+                        related_model = related_field.model
+                    else:
+                        field_object = related_model._meta.fields_map.get(related_field_name)
+
+                        if not field_object:
+                            raise FieldError(
+                                f"Unknown field {related_field_name} for model {related_model.__name__}"
+                            )
+                        related_table_field = related_table[
+                            field_object.source_field or related_field_name
+                        ]
+                        if func := field_object.get_for_dialect(
+                            related_model._meta.db.capabilities.dialect, "function_cast"
+                        ):
+                            related_table_field = func(field_object, related_table_field)
+                        distinct_on_by_source_field.append(related_table_field)
+            self.query.distinct_on(*distinct_on_by_source_field)
+
     def _resolve_annotate(self, fields_for_select: Collection[str] | None = None) -> bool:
         if not self._annotations:
             return False
@@ -405,7 +462,7 @@ def _clone(self) -> QuerySet[MODEL]:
         queryset._prefetch_queries = copy(self._prefetch_queries)
         queryset._single = self._single
         queryset._raise_does_not_exist = self._raise_does_not_exist
-        queryset._db = self._db
+        queryset._apply_db(self._db)
         queryset._limit = self._limit
         queryset._offset = self._offset
         queryset._fields_for_select = self._fields_for_select
@@ -627,16 +684,10 @@ def distinct(self, *args: str) -> QuerySet[MODEL]:
 
         :param args: Field names for ``DISTINCT ON`` (PostgreSQL only). Omit for plain
             ``DISTINCT``.
-        :raises OperationalError: If field arguments are given on a non-PostgreSQL database,
-            or if ``ORDER BY`` is specified but does not start with the ``DISTINCT ON`` fields.
         """
         queryset = self._clone()
         queryset._distinct = True
-        if args:
-            if isinstance(self.query, PostgreSQLQueryBuilder):
-                queryset._distinct_on = list(args)
-            else:
-                raise OperationalError("DISTINCT ON is only supported by PostgreSQL")
+        queryset._distinct_on = list(args)
         return queryset
 
     def union(self, *other_qs: QuerySet[Model], all: bool = False) -> UnionQuery[MODEL]:
@@ -1314,25 +1365,16 @@ def _make_query(self) -> None:
             self._fields_for_select,
         )
         self.resolve_filters()
+        self.resolve_distinct(
+            self._distinct,
+            self._distinct_on,
+            self._orderings,
+            self._annotations,
+        )
         if self._limit is not None:
             self.query._limit = self.query._wrapper_cls(self._limit)
         if self._offset is not None:
             self.query._offset = self.query._wrapper_cls(self._offset)
-        if self._distinct:
-            self.query._distinct = True
-            if isinstance(self.query, PostgreSQLQueryBuilder) and self._distinct_on:
-                ordering_fields = [ordering[0] for ordering in self._orderings]
-                len_ordering_fields = len(ordering_fields)
-                for i, field in enumerate(self._distinct_on):
-                    if ordering_fields and (
-                        i >= len_ordering_fields or ordering_fields[i] != field
-                    ):
-                        raise OperationalError(
-                            f"DISTINCT ON fields must match the leading ORDER BY fields. "
-                            f"Expected ORDER BY to start with {self._distinct_on!r}."
-                        )
-                self.query._distinct_on = []
-                self.query.distinct_on(*self._distinct_on)
         if self._select_for_update:
             self.query = self.query.for_update(
                 self._select_for_update_nowait,
@@ -1828,25 +1870,16 @@ def _make_query(self) -> None:
             fields_for_select=self._fields_for_select_list,
         )
         self.resolve_filters(self._fields_to_select_sql)
+        self.resolve_distinct(
+            self._distinct,
+            self._distinct_on,
+            self._orderings,
+            self._annotations,
+        )
         if self._limit:
             self.query._limit = self.query._wrapper_cls(self._limit)
         if self._offset:
             self.query._offset = self.query._wrapper_cls(self._offset)
-        if self._distinct:
-            self.query._distinct = True
-            if isinstance(self.query, PostgreSQLQueryBuilder) and self._distinct_on:
-                ordering_fields = [ordering[0] for ordering in self._orderings]
-                len_ordering_fields = len(ordering_fields)
-                for i, field in enumerate(self._distinct_on):
-                    if ordering_fields and (
-                        i >= len_ordering_fields or ordering_fields[i] != field
-                    ):
-                        raise OperationalError(
-                            f"DISTINCT ON fields must match the leading ORDER BY fields. "
-                            f"Expected ORDER BY to start with {self._distinct_on!r}."
-                        )
-                self.query._distinct_on = []
-                self.query.distinct_on(*self._distinct_on)
         if self._group_bys:
             self.query._groupbys = self._resolve_group_bys(*self._group_bys)
 
@@ -1966,6 +1999,12 @@ def _make_query(self) -> None:
             fields_for_select=self._fields_for_select.keys(),
         )
         self.resolve_filters()
+        self.resolve_distinct(
+            self._distinct,
+            self._distinct_on,
+            self._orderings,
+            self._annotations,
+        )
 
         # remove annotations that are not in fields_for_select
         self.query._selects = [
@@ -1976,21 +2015,6 @@ def _make_query(self) -> None:
             self.query._limit = self.query._wrapper_cls(self._limit)
         if self._offset:
             self.query._offset = self.query._wrapper_cls(self._offset)
-        if self._distinct:
-            self.query._distinct = True
-            if isinstance(self.query, PostgreSQLQueryBuilder) and self._distinct_on:
-                ordering_fields = [ordering[0] for ordering in self._orderings]
-                len_ordering_fields = len(ordering_fields)
-                for i, field in enumerate(self._distinct_on):
-                    if ordering_fields and (
-                        i >= len_ordering_fields or ordering_fields[i] != field
-                    ):
-                        raise OperationalError(
-                            f"DISTINCT ON fields must match the leading ORDER BY fields. "
-                            f"Expected ORDER BY to start with {self._distinct_on!r}."
-                        )
-                self.query._distinct_on = []
-                self.query.distinct_on(*self._distinct_on)
         if self._group_bys:
             self.query._groupbys = self._resolve_group_bys(*self._group_bys)
 
@@ -2521,7 +2545,7 @@ def _clone(self) -> UnionQuery[MODEL]:
         union._models = self._models
         union._union_query = None
         union._selects = self._selects
-        union._db = self._db
+        union._apply_db(self._db)
         union._qs = self._qs
         union._all = self._all
         union._orderings = self._orderings