Fixed #28586 -- Added model field fetch modes.

May your database queries be much reduced with minimal effort.

co-authored-by: Andreas Pelme <andreas@pelme.se>
co-authored-by: Simon Charette <charette.s@gmail.com>
co-authored-by: Jacob Walls <jacobtylerwalls@gmail.com>

7 files changed, 212 insertions, 6 deletions

diff --git a/tests/basic/tests.py b/tests/basic/tests.py
index 38d7d2a3d6..89aef16aef 100644
--- a/tests/basic/tests.py
+++ b/tests/basic/tests.py
@@ -807,6 +807,7 @@ class ManagerTest(SimpleTestCase):
         "alatest",
         "aupdate",
         "aupdate_or_create",
+        "fetch_mode",
     ]
 
     def test_manager_methods(self):
diff --git a/tests/defer/tests.py b/tests/defer/tests.py
index c0968080b1..29c63c566a 100644
--- a/tests/defer/tests.py
+++ b/tests/defer/tests.py
@@ -1,4 +1,5 @@
-from django.core.exceptions import FieldDoesNotExist, FieldError
+from django.core.exceptions import FieldDoesNotExist, FieldError, FieldFetchBlocked
+from django.db.models import FETCH_PEERS, RAISE
 from django.test import SimpleTestCase, TestCase
 
 from .models import (
@@ -29,6 +30,7 @@ class DeferTests(AssertionMixin, TestCase):
     def setUpTestData(cls):
         cls.s1 = Secondary.objects.create(first="x1", second="y1")
         cls.p1 = Primary.objects.create(name="p1", value="xx", related=cls.s1)
+        cls.p2 = Primary.objects.create(name="p2", value="yy", related=cls.s1)
 
     def test_defer(self):
         qs = Primary.objects.all()
@@ -141,7 +143,6 @@ class DeferTests(AssertionMixin, TestCase):
     def test_saving_object_with_deferred_field(self):
         # Saving models with deferred fields is possible (but inefficient,
         # since every field has to be retrieved first).
-        Primary.objects.create(name="p2", value="xy", related=self.s1)
         obj = Primary.objects.defer("value").get(name="p2")
         obj.name = "a new name"
         obj.save()
@@ -181,10 +182,71 @@ class DeferTests(AssertionMixin, TestCase):
         self.assertEqual(obj.name, "adonis")
 
     def test_defer_fk_attname(self):
-        primary = Primary.objects.defer("related_id").get()
+        primary = Primary.objects.defer("related_id").get(name="p1")
         with self.assertNumQueries(1):
             self.assertEqual(primary.related_id, self.p1.related_id)
 
+    def test_only_fetch_mode_fetch_peers(self):
+        p1, p2 = Primary.objects.fetch_mode(FETCH_PEERS).only("name")
+        with self.assertNumQueries(1):
+            p1.value
+        with self.assertNumQueries(0):
+            p2.value
+
+    def test_only_fetch_mode_fetch_peers_single(self):
+        p1 = Primary.objects.fetch_mode(FETCH_PEERS).only("name").get(name="p1")
+        with self.assertNumQueries(1):
+            p1.value
+
+    def test_defer_fetch_mode_fetch_peers(self):
+        p1, p2 = Primary.objects.fetch_mode(FETCH_PEERS).defer("value")
+        with self.assertNumQueries(1):
+            p1.value
+        with self.assertNumQueries(0):
+            p2.value
+
+    def test_defer_fetch_mode_fetch_peers_single(self):
+        p1 = Primary.objects.fetch_mode(FETCH_PEERS).defer("value").get(name="p1")
+        with self.assertNumQueries(1):
+            p1.value
+
+    def test_only_fetch_mode_raise(self):
+        p1 = Primary.objects.fetch_mode(RAISE).only("name").get(name="p1")
+        msg = "Fetching of Primary.value blocked."
+        with self.assertRaisesMessage(FieldFetchBlocked, msg) as cm:
+            p1.value
+        self.assertIsNone(cm.exception.__cause__)
+        self.assertTrue(cm.exception.__suppress_context__)
+
+    def test_defer_fetch_mode_raise(self):
+        p1 = Primary.objects.fetch_mode(RAISE).defer("value").get(name="p1")
+        msg = "Fetching of Primary.value blocked."
+        with self.assertRaisesMessage(FieldFetchBlocked, msg) as cm:
+            p1.value
+        self.assertIsNone(cm.exception.__cause__)
+        self.assertTrue(cm.exception.__suppress_context__)
+
+
+class DeferOtherDatabaseTests(TestCase):
+    databases = {"other"}
+
+    @classmethod
+    def setUpTestData(cls):
+        cls.s1 = Secondary.objects.using("other").create(first="x1", second="y1")
+        cls.p1 = Primary.objects.using("other").create(
+            name="p1", value="xx", related=cls.s1
+        )
+        cls.p2 = Primary.objects.using("other").create(
+            name="p2", value="yy", related=cls.s1
+        )
+
+    def test_defer_fetch_mode_fetch_peers(self):
+        p1, p2 = Primary.objects.using("other").fetch_mode(FETCH_PEERS).defer("value")
+        with self.assertNumQueries(1, using="other"):
+            p1.value
+        with self.assertNumQueries(0, using="other"):
+            p2.value
+
 
 class BigChildDeferTests(AssertionMixin, TestCase):
     @classmethod
diff --git a/tests/generic_relations/tests.py b/tests/generic_relations/tests.py
index 1b53dbd8f4..3de243d7b8 100644
--- a/tests/generic_relations/tests.py
+++ b/tests/generic_relations/tests.py
@@ -1,7 +1,8 @@
 from django.contrib.contenttypes.models import ContentType
 from django.contrib.contenttypes.prefetch import GenericPrefetch
-from django.core.exceptions import FieldError
+from django.core.exceptions import FieldError, FieldFetchBlocked
 from django.db.models import Q, prefetch_related_objects
+from django.db.models.fetch_modes import FETCH_PEERS, RAISE
 from django.test import SimpleTestCase, TestCase, skipUnlessDBFeature
 
 from .models import (
@@ -780,6 +781,46 @@ class GenericRelationsTests(TestCase):
                 self.platypus.latin_name,
             )
 
+    def test_fetch_mode_fetch_peers(self):
+        TaggedItem.objects.bulk_create(
+            [
+                TaggedItem(tag="lion", content_object=self.lion),
+                TaggedItem(tag="platypus", content_object=self.platypus),
+                TaggedItem(tag="quartz", content_object=self.quartz),
+            ]
+        )
+        # Peers fetching should fetch all related peers GFKs at once which is
+        # one query per content type.
+        with self.assertNumQueries(1):
+            quartz_tag, platypus_tag, lion_tag = TaggedItem.objects.fetch_mode(
+                FETCH_PEERS
+            ).order_by("-pk")[:3]
+        with self.assertNumQueries(2):
+            self.assertEqual(lion_tag.content_object, self.lion)
+        with self.assertNumQueries(0):
+            self.assertEqual(platypus_tag.content_object, self.platypus)
+            self.assertEqual(quartz_tag.content_object, self.quartz)
+        # It should ignore already cached instances though.
+        with self.assertNumQueries(1):
+            quartz_tag, platypus_tag, lion_tag = TaggedItem.objects.fetch_mode(
+                FETCH_PEERS
+            ).order_by("-pk")[:3]
+        with self.assertNumQueries(2):
+            self.assertEqual(quartz_tag.content_object, self.quartz)
+            self.assertEqual(lion_tag.content_object, self.lion)
+        with self.assertNumQueries(0):
+            self.assertEqual(platypus_tag.content_object, self.platypus)
+            self.assertEqual(quartz_tag.content_object, self.quartz)
+
+    def test_fetch_mode_raise(self):
+        TaggedItem.objects.create(tag="lion", content_object=self.lion)
+        tag = TaggedItem.objects.fetch_mode(RAISE).get(tag="yellow")
+        msg = "Fetching of TaggedItem.content_object blocked."
+        with self.assertRaisesMessage(FieldFetchBlocked, msg) as cm:
+            tag.content_object
+        self.assertIsNone(cm.exception.__cause__)
+        self.assertTrue(cm.exception.__suppress_context__)
+
 
 class ProxyRelatedModelTest(TestCase):
     def test_default_behavior(self):
diff --git a/tests/many_to_one/tests.py b/tests/many_to_one/tests.py
index ac43c0da95..c5fa458570 100644
--- a/tests/many_to_one/tests.py
+++ b/tests/many_to_one/tests.py
@@ -1,8 +1,13 @@
 import datetime
 from copy import deepcopy
 
-from django.core.exceptions import FieldError, MultipleObjectsReturned
+from django.core.exceptions import (
+    FieldError,
+    FieldFetchBlocked,
+    MultipleObjectsReturned,
+)
 from django.db import IntegrityError, models, transaction
+from django.db.models import FETCH_PEERS, RAISE
 from django.test import TestCase
 from django.utils.translation import gettext_lazy
 
@@ -916,3 +921,23 @@ class ManyToOneTests(TestCase):
                 instances=countries,
                 querysets=[City.objects.all(), City.objects.all()],
             )
+
+    def test_fetch_mode_fetch_peers_forward(self):
+        Article.objects.create(
+            headline="This is another test",
+            pub_date=datetime.date(2005, 7, 27),
+            reporter=self.r2,
+        )
+        a1, a2 = Article.objects.fetch_mode(FETCH_PEERS)
+        with self.assertNumQueries(1):
+            a1.reporter
+        with self.assertNumQueries(0):
+            a2.reporter
+
+    def test_fetch_mode_raise_forward(self):
+        a = Article.objects.fetch_mode(RAISE).get(pk=self.a.pk)
+        msg = "Fetching of Article.reporter blocked."
+        with self.assertRaisesMessage(FieldFetchBlocked, msg) as cm:
+            a.reporter
+        self.assertIsNone(cm.exception.__cause__)
+        self.assertTrue(cm.exception.__suppress_context__)
diff --git a/tests/one_to_one/tests.py b/tests/one_to_one/tests.py
index d9bcb5d4dc..da7bd992c0 100644
--- a/tests/one_to_one/tests.py
+++ b/tests/one_to_one/tests.py
@@ -1,4 +1,6 @@
+from django.core.exceptions import FieldFetchBlocked
 from django.db import IntegrityError, connection, transaction
+from django.db.models import FETCH_PEERS, RAISE
 from django.test import TestCase
 
 from .models import (
@@ -619,3 +621,39 @@ class OneToOneTests(TestCase):
                 instances=places,
                 querysets=[Bar.objects.all(), Bar.objects.all()],
             )
+
+    def test_fetch_mode_fetch_peers_forward(self):
+        Restaurant.objects.create(
+            place=self.p2, serves_hot_dogs=True, serves_pizza=False
+        )
+        r1, r2 = Restaurant.objects.fetch_mode(FETCH_PEERS)
+        with self.assertNumQueries(1):
+            r1.place
+        with self.assertNumQueries(0):
+            r2.place
+
+    def test_fetch_mode_fetch_peers_reverse(self):
+        Restaurant.objects.create(
+            place=self.p2, serves_hot_dogs=True, serves_pizza=False
+        )
+        p1, p2 = Place.objects.fetch_mode(FETCH_PEERS)
+        with self.assertNumQueries(1):
+            p1.restaurant
+        with self.assertNumQueries(0):
+            p2.restaurant
+
+    def test_fetch_mode_raise_forward(self):
+        r = Restaurant.objects.fetch_mode(RAISE).get(pk=self.r1.pk)
+        msg = "Fetching of Restaurant.place blocked."
+        with self.assertRaisesMessage(FieldFetchBlocked, msg) as cm:
+            r.place
+        self.assertIsNone(cm.exception.__cause__)
+        self.assertTrue(cm.exception.__suppress_context__)
+
+    def test_fetch_mode_raise_reverse(self):
+        p = Place.objects.fetch_mode(RAISE).get(pk=self.p1.pk)
+        msg = "Fetching of Place.restaurant blocked."
+        with self.assertRaisesMessage(FieldFetchBlocked, msg) as cm:
+            p.restaurant
+        self.assertIsNone(cm.exception.__cause__)
+        self.assertTrue(cm.exception.__suppress_context__)
diff --git a/tests/prefetch_related/tests.py b/tests/prefetch_related/tests.py
index 54b197ad83..6e4acdddf6 100644
--- a/tests/prefetch_related/tests.py
+++ b/tests/prefetch_related/tests.py
@@ -4,6 +4,7 @@ from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import ObjectDoesNotExist
 from django.db import NotSupportedError, connection
 from django.db.models import F, Prefetch, QuerySet, prefetch_related_objects
+from django.db.models.fetch_modes import RAISE
 from django.db.models.query import get_prefetcher
 from django.db.models.sql import Query
 from django.test import (
@@ -107,6 +108,10 @@ class PrefetchRelatedTests(TestDataMixin, TestCase):
         normal_books = [a.first_book for a in Author.objects.all()]
         self.assertEqual(books, normal_books)
 
+    def test_fetch_mode_raise(self):
+        authors = list(Author.objects.fetch_mode(RAISE).prefetch_related("first_book"))
+        authors[0].first_book  # No exception, already loaded
+
     def test_foreignkey_reverse(self):
         with self.assertNumQueries(2):
             [
diff --git a/tests/raw_query/tests.py b/tests/raw_query/tests.py
index 853b7ee20e..f66afbf28b 100644
--- a/tests/raw_query/tests.py
+++ b/tests/raw_query/tests.py
@@ -1,7 +1,8 @@
 from datetime import date
 from decimal import Decimal
 
-from django.core.exceptions import FieldDoesNotExist
+from django.core.exceptions import FieldDoesNotExist, FieldFetchBlocked
+from django.db.models import FETCH_PEERS, RAISE
 from django.db.models.query import RawQuerySet
 from django.test import TestCase, skipUnlessDBFeature
 
@@ -158,6 +159,22 @@ class RawQueryTests(TestCase):
         books = Book.objects.all()
         self.assertSuccessfulRawQuery(Book, query, books)
 
+    def test_fk_fetch_mode_peers(self):
+        query = "SELECT * FROM raw_query_book"
+        books = list(Book.objects.fetch_mode(FETCH_PEERS).raw(query))
+        with self.assertNumQueries(1):
+            books[0].author
+            books[1].author
+
+    def test_fk_fetch_mode_raise(self):
+        query = "SELECT * FROM raw_query_book"
+        books = list(Book.objects.fetch_mode(RAISE).raw(query))
+        msg = "Fetching of Book.author blocked."
+        with self.assertRaisesMessage(FieldFetchBlocked, msg) as cm:
+            books[0].author
+        self.assertIsNone(cm.exception.__cause__)
+        self.assertTrue(cm.exception.__suppress_context__)
+
     def test_db_column_handler(self):
         """
         Test of a simple raw query against a model containing a field with
@@ -294,6 +311,23 @@ class RawQueryTests(TestCase):
         with self.assertRaisesMessage(FieldDoesNotExist, msg):
             list(Author.objects.raw(query))
 
+    def test_missing_fields_fetch_mode_peers(self):
+        query = "SELECT id, first_name, dob FROM raw_query_author"
+        authors = list(Author.objects.fetch_mode(FETCH_PEERS).raw(query))
+        with self.assertNumQueries(1):
+            authors[0].last_name
+            authors[1].last_name
+
+    def test_missing_fields_fetch_mode_raise(self):
+        query = "SELECT id, first_name, dob FROM raw_query_author"
+        authors = list(Author.objects.fetch_mode(RAISE).raw(query))
+        msg = "Fetching of Author.last_name blocked."
+        with self.assertRaisesMessage(FieldFetchBlocked, msg) as cm:
+            authors[0].last_name
+        self.assertIsNone(cm.exception.__cause__)
+        self.assertTrue(cm.exception.__suppress_context__)
+        self.assertTrue(cm.exception.__suppress_context__)
+
     def test_annotations(self):
         query = (
             "SELECT a.*, count(b.id) as book_count "