Support relations in .only() (#1923)

henadzit · web-flow · commit 21463c746110 · 2025-03-27T16:44:01.000+08:00
* Make expand_field_expression reusable

* .only() supports relations

* Remove redundant check

* Prohibit using .values() and .values_list() with .only()

* Add tests

* Add a comment

* Update CHANGELOG

* CR feedback
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -6,14 +6,23 @@ Changelog
 
 .. rst-class:: emphasize-children
 
-0.24
+0.25
 ====
 
-0.24.3 (unreleased)
+0.25.0 (unreleased)
 ------
 Changed
 ^^^^^^^
 - Skip database selection if the router is not configured to improve performance (#1915)
+- `.values()`, `.values_list()` and `.only()` cannot be used together (#1923)
+
+Added
+^^^^^
+- `.only` supports selecting related fields, e.g. `.only("related__field")` (#1923)
+
+
+0.24
+====
 
 0.24.2
 ------
diff --git a/tests/test_aggregation.py b/tests/test_aggregation.py
@@ -11,7 +11,7 @@
 )
 from tortoise.contrib import test
 from tortoise.contrib.test.condition import In
-from tortoise.exceptions import ConfigurationError
+from tortoise.exceptions import FieldError
 from tortoise.expressions import F, Q
 from tortoise.functions import Avg, Coalesce, Concat, Count, Lower, Max, Min, Sum, Trim
 
@@ -60,7 +60,7 @@ async def test_aggregation(self):
             event_with_annotation.tournament_id,
         )
 
-        with self.assertRaisesRegex(ConfigurationError, "name__id not resolvable"):
+        with self.assertRaisesRegex(FieldError, "name__id not resolvable"):
             await Event.all().annotate(tournament_test_id=Sum("name__id")).first()
 
     async def test_nested_aggregation_in_annotation(self):
diff --git a/tests/test_only.py b/tests/test_only.py
@@ -1,6 +1,7 @@
-from tests.testmodels import SourceFields, StraightFields
+from tests.testmodels import DoubleFK, Event, SourceFields, StraightFields, Tournament
 from tortoise.contrib import test
-from tortoise.exceptions import IncompleteInstanceError
+from tortoise.functions import Count
+from tortoise.exceptions import FieldError, IncompleteInstanceError
 
 
 class TestOnlyStraight(test.TestCase):
@@ -64,3 +65,230 @@ async def asyncSetUp(self) -> None:
         await super().asyncSetUp()
         self.model = SourceFields  # type: ignore
         self.instance = await self.model.create(chars="Test")
+
+
+class TestOnlyRecursive(test.TestCase):
+    async def test_one_level(self):
+        left_1st_lvl = await DoubleFK.create(name="1st")
+        root = await DoubleFK.create(name="root", left=left_1st_lvl)
+
+        ret = (
+            await DoubleFK.filter(pk=root.pk).only("name", "left__name", "left__left__name").first()
+        )
+        self.assertIsNotNone(ret)
+        with self.assertRaises(AttributeError):
+            _ = ret.id
+        self.assertEqual(ret.name, "root")
+        self.assertEqual(ret.left.name, "1st")
+        with self.assertRaises(AttributeError):
+            _ = ret.left.id
+        with self.assertRaises(AttributeError):
+            _ = ret.right
+
+    async def test_two_levels(self):
+        left_2nd_lvl = await DoubleFK.create(name="second leaf")
+        left_1st_lvl = await DoubleFK.create(name="1st", left=left_2nd_lvl)
+        root = await DoubleFK.create(name="root", left=left_1st_lvl)
+
+        ret = (
+            await DoubleFK.filter(pk=root.pk).only("name", "left__name", "left__left__name").first()
+        )
+        self.assertIsNotNone(ret)
+        with self.assertRaises(AttributeError):
+            _ = ret.id
+        self.assertEqual(ret.name, "root")
+        self.assertEqual(ret.left.name, "1st")
+        with self.assertRaises(AttributeError):
+            _ = ret.left.id
+        self.assertEqual(ret.left.left.name, "second leaf")
+
+    async def test_two_levels_reverse_argument_order(self):
+        left_2nd_lvl = await DoubleFK.create(name="second leaf")
+        left_1st_lvl = await DoubleFK.create(name="1st", left=left_2nd_lvl)
+        root = await DoubleFK.create(name="root", left=left_1st_lvl)
+
+        ret = (
+            await DoubleFK.filter(pk=root.pk).only("left__left__name", "left__name", "name").first()
+        )
+        self.assertIsNotNone(ret)
+        with self.assertRaises(AttributeError):
+            _ = ret.id
+        self.assertEqual(ret.name, "root")
+        self.assertEqual(ret.left.name, "1st")
+        with self.assertRaises(AttributeError):
+            _ = ret.left.id
+        self.assertEqual(ret.left.left.name, "second leaf")
+
+
+class TestOnlyRelated(test.TestCase):
+    async def test_related_one_level(self):
+        tournament = await Tournament.create(name="New Tournament", desc="New Description")
+        await Event.create(name="Event 1", tournament=tournament)
+        await Event.create(name="Event 2", tournament=tournament)
+
+        ret = (
+            await Event.filter(tournament=tournament)
+            .only("name", "tournament__name")
+            .order_by("name")
+        )
+        self.assertEqual(len(ret), 2)
+        self.assertEqual(ret[0].name, "Event 1")
+        with self.assertRaises(AttributeError):
+            _ = ret[0].alias
+        self.assertEqual(ret[1].name, "Event 2")
+        with self.assertRaises(AttributeError):
+            _ = ret[1].alias
+        self.assertEqual(ret[0].tournament.name, "New Tournament")
+        with self.assertRaises(AttributeError):
+            _ = ret[0].tournament.id
+        with self.assertRaises(AttributeError):
+            _ = ret[0].tournament.desc
+
+    async def test_related_one_level_reversed_argument_order(self):
+        tournament = await Tournament.create(name="New Tournament", desc="New Description")
+        await Event.create(name="Event 1", tournament=tournament)
+        await Event.create(name="Event 2", tournament=tournament)
+
+        ret = (
+            await Event.filter(tournament=tournament)
+            .only("tournament__name", "name")
+            .order_by("name")
+        )
+        self.assertEqual(len(ret), 2)
+        self.assertEqual(ret[0].name, "Event 1")
+        self.assertEqual(ret[0].tournament.name, "New Tournament")
+
+    async def test_just_related(self):
+        tournament = await Tournament.create(name="New Tournament", desc="New Description")
+        await Event.create(name="Event 1", tournament=tournament)
+        await Event.create(name="Event 2", tournament=tournament)
+
+        ret = (
+            await Event.filter(tournament=tournament)
+            .only("tournament__name")
+            .order_by("name")
+            .all()
+        )
+        self.assertEqual(len(ret), 2)
+        self.assertEqual(ret[0].tournament.name, "New Tournament")
+        self.assertEqual(ret[1].tournament.name, "New Tournament")
+
+
+class TestOnlyAdvanced(test.TestCase):
+    async def asyncSetUp(self) -> None:
+        await super().asyncSetUp()
+        self.tournament = await Tournament.create(name="Tournament A", desc="Description A")
+        self.event1 = await Event.create(name="Event 1", tournament=self.tournament)
+        self.event2 = await Event.create(name="Event 2", tournament=self.tournament)
+
+    async def test_exclude(self):
+        """Test .only() combined with .exclude()"""
+        events = await Event.filter(tournament=self.tournament).exclude(name="Event 2").only("name")
+        self.assertEqual(len(events), 1)
+        self.assertEqual(events[0].name, "Event 1")
+        with self.assertRaises(AttributeError):
+            _ = events[0].modified
+
+    async def test_limit(self):
+        """Test .only() combined with .limit()"""
+        events = await Event.all().only("name").limit(1)
+        self.assertEqual(len(events), 1)
+        self.assertEqual(events[0].name, "Event 1")  # Assumes ordering by PK
+        with self.assertRaises(AttributeError):
+            _ = events[0].modified
+
+    async def test_distinct(self):
+        """Test .only() combined with .distinct()"""
+        # Create duplicate event names
+        await Event.create(name="Event 1", tournament=self.tournament)
+
+        events = await Event.all().only("name").distinct()
+        # Should only have 2 distinct event names
+        self.assertEqual(len(events), 2)
+        event_names = {e.name for e in events}
+        self.assertEqual(event_names, {"Event 1", "Event 2"})
+
+    async def test_values(self):
+        """Test .only() combined with .values()"""
+        with self.assertRaises(ValueError, msg="values() cannot be used with .only()"):
+            await Event.all().only("name").values("name")
+
+    async def test_pk_field(self):
+        """Test .only() with just the primary key field"""
+        tournament = await Tournament.first().only("id")
+        self.assertIsNotNone(tournament.id)
+        with self.assertRaises(AttributeError):
+            _ = tournament.name
+
+    async def test_empty(self):
+        """Test .only() with no fields (should raise an error)"""
+        with self.assertRaises(ValueError):
+            await Event.all().only()
+
+    async def test_annotate(self):
+        tournaments = await Tournament.annotate(event_count=Count("events")).only(
+            "name", "event_count"
+        )
+
+        self.assertEqual(tournaments[0].name, "Tournament A")
+        self.assertEqual(tournaments[0].event_count, 2)
+        with self.assertRaises(AttributeError):
+            _ = tournaments[0].desc
+
+    async def test_nonexistent_field(self):
+        """Test .only() with a field that doesn't exist"""
+        with self.assertRaises(FieldError):
+            await Event.all().only("nonexistent_field").all()
+
+    async def test_join_in_filter(self):
+        event = await Event.filter(name="Event 1").only("name").first()
+        self.assertEqual(event.name, "Event 1")
+        with self.assertRaises(AttributeError):
+            _ = event.tournament
+
+        event = await Event.filter(tournament__name="Tournament A").only("name").first()
+        self.assertEqual(event.name, "Event 1")
+        with self.assertRaises(AttributeError):
+            _ = event.tournament
+
+        event = (
+            await Event.filter(tournament__name="Tournament A")
+            .only("name", "tournament__name")
+            .first()
+        )
+        self.assertEqual(event.name, "Event 1")
+        self.assertEqual(event.tournament.name, "Tournament A")
+
+    async def test_join_in_order_by(self):
+        events = await Event.all().order_by("name").only("name")
+        self.assertEqual(events[0].name, "Event 1")
+        with self.assertRaises(AttributeError):
+            _ = events[0].tournament
+
+        events = await Event.all().order_by("tournament__name", "name").only("name")
+        self.assertEqual(events[0].name, "Event 1")
+        with self.assertRaises(AttributeError):
+            _ = events[0].tournament
+
+        events = (
+            await Event.all().order_by("tournament__name", "name").only("name", "tournament__name")
+        )
+        self.assertEqual(events[0].name, "Event 1")
+        self.assertEqual(events[0].tournament.name, "Tournament A")
+
+    async def test_select_related(self):
+        """Test .only() with .select_related() for basic functionality"""
+        event = (
+            await Event.filter(name="Event 1")
+            .select_related("tournament")
+            .only("name", "tournament__name")
+            .first()
+        )
+
+        self.assertEqual(event.name, "Event 1")
+        self.assertEqual(event.tournament.name, "Tournament A")
+
+        with self.assertRaises(AttributeError):
+            _ = event.id
+        with self.assertRaises(AttributeError):
+            _ = event.tournament.id
diff --git a/tortoise/backends/base/executor.py b/tortoise/backends/base/executor.py
@@ -114,10 +114,15 @@ async def execute_select(
                 for model, index, *__, full_path in self.select_related_idx[1:]:
                     (*path, attr) = full_path
                     related_items = row_items[current_idx : current_idx + index]
-                    if not any((v for _, v in related_items)):
-                        obj = None
-                    else:
+                    if any(v for _, v in related_items):
                         obj = model._init_from_db(**{k.split(".")[1]: v for k, v in related_items})
+                    elif index == 0:
+                        # 0 signals that an empty "filler" object should be created in the case
+                        # where a field of related model is selected but model itself isn't,
+                        # e.g. .only("relatedmodel__field")
+                        obj = model._init_from_db()
+                    else:
+                        obj = None
                     target = instances.get(tuple(path))
                     if target is not None:
                         setattr(target, f"_{attr}", obj)
diff --git a/tortoise/models.py b/tortoise/models.py
@@ -1066,7 +1066,9 @@ async def refresh_from_db(
         if not self._saved_in_db:
             raise OperationalError("Can't refresh unpersisted record")
         db = using_db or self._choose_db()
-        qs = QuerySet(self.__class__).using_db(db).only(*(fields or []))
+        qs = QuerySet(self.__class__).using_db(db)
+        if fields:
+            qs = qs.only(*fields)
         obj = await qs.get(pk=self.pk)
 
         for field in fields or self._meta.db_fields:
diff --git a/tortoise/query_utils.py b/tortoise/query_utils.py
diff --git a/tortoise/queryset.py b/tortoise/queryset.py

Original file line number	Diff line number	Diff line change
`@@ -11,7 +11,7 @@`
`11`	`11`	`)`
`12`	`12`	`from tortoise.contrib import test`
`13`	`13`	`from tortoise.contrib.test.condition import In`
`14`		`-from tortoise.exceptions import ConfigurationError`
	`14`	`+from tortoise.exceptions import FieldError`
`15`	`15`	`from tortoise.expressions import F, Q`
`16`	`16`	`from tortoise.functions import Avg, Coalesce, Concat, Count, Lower, Max, Min, Sum, Trim`
`17`	`17`
`@@ -60,7 +60,7 @@ async def test_aggregation(self):`
`60`	`60`	`event_with_annotation.tournament_id,`
`61`	`61`	`)`
`62`	`62`
`63`		`- with self.assertRaisesRegex(ConfigurationError, "name__id not resolvable"):`
	`63`	`+ with self.assertRaisesRegex(FieldError, "name__id not resolvable"):`
`64`	`64`	`await Event.all().annotate(tournament_test_id=Sum("name__id")).first()`
`65`	`65`
`66`	`66`	`async def test_nested_aggregation_in_annotation(self):`