Fixed #33651 -- Added support for prefetching GenericForeignKey.

Co-authored-by: revanthgss <revanthgss@almabase.com>
Co-authored-by: Mariusz Felisiak <felisiak.mariusz@gmail.com>

7 files changed, 210 insertions, 10 deletions

diff --git a/tests/contenttypes_tests/test_fields.py b/tests/contenttypes_tests/test_fields.py
index 418669140b..5510f34cd0 100644
--- a/tests/contenttypes_tests/test_fields.py
+++ b/tests/contenttypes_tests/test_fields.py
@@ -1,9 +1,11 @@
 import json
 
 from django.contrib.contenttypes.fields import GenericForeignKey
+from django.contrib.contenttypes.prefetch import GenericPrefetch
 from django.db import models
 from django.test import TestCase
 from django.test.utils import isolate_apps
+from django.utils.deprecation import RemovedInDjango60Warning
 
 from .models import Answer, Post, Question
 
@@ -22,14 +24,6 @@ class GenericForeignKeyTests(TestCase):
         ):
             Answer.question.get_content_type()
 
-    def test_incorrect_get_prefetch_queryset_arguments(self):
-        with self.assertRaisesMessage(
-            ValueError, "Custom queryset can't be used for this lookup."
-        ):
-            Answer.question.get_prefetch_queryset(
-                Answer.objects.all(), Answer.objects.all()
-            )
-
     def test_get_object_cache_respects_deleted_objects(self):
         question = Question.objects.create(text="Who?")
         post = Post.objects.create(title="Answer", parent=question)
@@ -59,3 +53,55 @@ class GenericRelationTests(TestCase):
         answer2 = Answer.objects.create(question=question)
         result = json.loads(Question.answer_set.field.value_to_string(question))
         self.assertCountEqual(result, [answer1.pk, answer2.pk])
+
+
+class GetPrefetchQuerySetDeprecation(TestCase):
+    def test_generic_relation_warning(self):
+        Question.objects.create(text="test")
+        questions = Question.objects.all()
+        msg = (
+            "get_prefetch_queryset() is deprecated. Use get_prefetch_querysets() "
+            "instead."
+        )
+        with self.assertWarnsMessage(RemovedInDjango60Warning, msg):
+            questions[0].answer_set.get_prefetch_queryset(questions)
+
+    def test_generic_foreign_key_warning(self):
+        answers = Answer.objects.all()
+        msg = (
+            "get_prefetch_queryset() is deprecated. Use get_prefetch_querysets() "
+            "instead."
+        )
+        with self.assertWarnsMessage(RemovedInDjango60Warning, msg):
+            Answer.question.get_prefetch_queryset(answers)
+
+
+class GetPrefetchQuerySetsTests(TestCase):
+    def test_duplicate_querysets(self):
+        question = Question.objects.create(text="What is your name?")
+        answer = Answer.objects.create(text="Joe", question=question)
+        answer = Answer.objects.get(pk=answer.pk)
+        msg = "Only one queryset is allowed for each content type."
+        with self.assertRaisesMessage(ValueError, msg):
+            models.prefetch_related_objects(
+                [answer],
+                GenericPrefetch(
+                    "question",
+                    [
+                        Question.objects.all(),
+                        Question.objects.filter(text__startswith="test"),
+                    ],
+                ),
+            )
+
+    def test_generic_relation_invalid_length(self):
+        Question.objects.create(text="test")
+        questions = Question.objects.all()
+        msg = (
+            "querysets argument of get_prefetch_querysets() should have a length of 1."
+        )
+        with self.assertRaisesMessage(ValueError, msg):
+            questions[0].answer_set.get_prefetch_querysets(
+                instances=questions,
+                querysets=[Answer.objects.all(), Question.objects.all()],
+            )
diff --git a/tests/contenttypes_tests/test_models.py b/tests/contenttypes_tests/test_models.py
index 95518bb3df..36c14cf56f 100644
--- a/tests/contenttypes_tests/test_models.py
+++ b/tests/contenttypes_tests/test_models.py
@@ -1,5 +1,6 @@
 from django.apps import apps
 from django.contrib.contenttypes.models import ContentType, ContentTypeManager
+from django.contrib.contenttypes.prefetch import GenericPrefetch
 from django.db import models
 from django.db.migrations.state import ProjectState
 from django.test import TestCase, override_settings
@@ -328,3 +329,17 @@ class ContentTypesMultidbTests(TestCase):
             1, using="other"
         ):
             ContentType.objects.get_for_model(Author)
+
+
+class GenericPrefetchTests(TestCase):
+    def test_values_queryset(self):
+        msg = "Prefetch querysets cannot use raw(), values(), and values_list()."
+        with self.assertRaisesMessage(ValueError, msg):
+            GenericPrefetch("question", [Author.objects.values("pk")])
+        with self.assertRaisesMessage(ValueError, msg):
+            GenericPrefetch("question", [Author.objects.values_list("pk")])
+
+    def test_raw_queryset(self):
+        msg = "Prefetch querysets cannot use raw(), values(), and values_list()."
+        with self.assertRaisesMessage(ValueError, msg):
+            GenericPrefetch("question", [Author.objects.raw("select pk from author")])
diff --git a/tests/generic_relations/tests.py b/tests/generic_relations/tests.py
index fab23dfde5..e0c6fe2db7 100644
--- a/tests/generic_relations/tests.py
+++ b/tests/generic_relations/tests.py
@@ -1,6 +1,7 @@
 from django.contrib.contenttypes.models import ContentType
+from django.contrib.contenttypes.prefetch import GenericPrefetch
 from django.core.exceptions import FieldError
-from django.db.models import Q
+from django.db.models import Q, prefetch_related_objects
 from django.test import SimpleTestCase, TestCase, skipUnlessDBFeature
 
 from .models import (
@@ -747,6 +748,38 @@ class GenericRelationsTests(TestCase):
                 comparison.first_obj.comparisons.all(), [comparison]
             )
 
+    def test_generic_prefetch(self):
+        tagged_vegetable = TaggedItem.objects.create(
+            tag="great", content_object=self.bacon
+        )
+        tagged_animal = TaggedItem.objects.create(
+            tag="awesome", content_object=self.platypus
+        )
+        # Getting the instances again so that content object is deferred.
+        tagged_vegetable = TaggedItem.objects.get(pk=tagged_vegetable.pk)
+        tagged_animal = TaggedItem.objects.get(pk=tagged_animal.pk)
+
+        with self.assertNumQueries(2):
+            prefetch_related_objects(
+                [tagged_vegetable, tagged_animal],
+                GenericPrefetch(
+                    "content_object",
+                    [Vegetable.objects.all(), Animal.objects.only("common_name")],
+                ),
+            )
+        with self.assertNumQueries(0):
+            self.assertEqual(tagged_vegetable.content_object.name, self.bacon.name)
+        with self.assertNumQueries(0):
+            self.assertEqual(
+                tagged_animal.content_object.common_name,
+                self.platypus.common_name,
+            )
+        with self.assertNumQueries(1):
+            self.assertEqual(
+                tagged_animal.content_object.latin_name,
+                self.platypus.latin_name,
+            )
+
 
 class ProxyRelatedModelTest(TestCase):
     def test_default_behavior(self):
diff --git a/tests/many_to_many/tests.py b/tests/many_to_many/tests.py
index 9b43d52c8a..7ed3b80abc 100644
--- a/tests/many_to_many/tests.py
+++ b/tests/many_to_many/tests.py
@@ -2,6 +2,7 @@ from unittest import mock
 
 from django.db import transaction
 from django.test import TestCase, skipIfDBFeature, skipUnlessDBFeature
+from django.utils.deprecation import RemovedInDjango60Warning
 
 from .models import Article, InheritedArticleA, InheritedArticleB, Publication, User
 
@@ -561,3 +562,23 @@ class ManyToManyTests(TestCase):
         self.assertEqual(
             self.p3.article_set.exists(), self.p3.article_set.all().exists()
         )
+
+    def test_get_prefetch_queryset_warning(self):
+        articles = Article.objects.all()
+        msg = (
+            "get_prefetch_queryset() is deprecated. Use get_prefetch_querysets() "
+            "instead."
+        )
+        with self.assertWarnsMessage(RemovedInDjango60Warning, msg):
+            self.a1.publications.get_prefetch_queryset(articles)
+
+    def test_get_prefetch_querysets_invalid_querysets_length(self):
+        articles = Article.objects.all()
+        msg = (
+            "querysets argument of get_prefetch_querysets() should have a length of 1."
+        )
+        with self.assertRaisesMessage(ValueError, msg):
+            self.a1.publications.get_prefetch_querysets(
+                instances=articles,
+                querysets=[Publication.objects.all(), Publication.objects.all()],
+            )
diff --git a/tests/many_to_one/tests.py b/tests/many_to_one/tests.py
index 7a6d112a09..b07972ec31 100644
--- a/tests/many_to_one/tests.py
+++ b/tests/many_to_one/tests.py
@@ -4,6 +4,7 @@ from copy import deepcopy
 from django.core.exceptions import FieldError, MultipleObjectsReturned
 from django.db import IntegrityError, models, transaction
 from django.test import TestCase
+from django.utils.deprecation import RemovedInDjango60Warning
 from django.utils.translation import gettext_lazy
 
 from .models import (
@@ -877,3 +878,49 @@ class ManyToOneTests(TestCase):
             usa.cities.remove(chicago.pk)
         with self.assertRaisesMessage(TypeError, msg):
             usa.cities.set([chicago.pk])
+
+    def test_get_prefetch_queryset_warning(self):
+        City.objects.create(name="Chicago")
+        cities = City.objects.all()
+        msg = (
+            "get_prefetch_queryset() is deprecated. Use get_prefetch_querysets() "
+            "instead."
+        )
+        with self.assertWarnsMessage(RemovedInDjango60Warning, msg):
+            City.country.get_prefetch_queryset(cities)
+
+    def test_get_prefetch_queryset_reverse_warning(self):
+        usa = Country.objects.create(name="United States")
+        City.objects.create(name="Chicago")
+        countries = Country.objects.all()
+        msg = (
+            "get_prefetch_queryset() is deprecated. Use get_prefetch_querysets() "
+            "instead."
+        )
+        with self.assertWarnsMessage(RemovedInDjango60Warning, msg):
+            usa.cities.get_prefetch_queryset(countries)
+
+    def test_get_prefetch_querysets_invalid_querysets_length(self):
+        City.objects.create(name="Chicago")
+        cities = City.objects.all()
+        msg = (
+            "querysets argument of get_prefetch_querysets() should have a length of 1."
+        )
+        with self.assertRaisesMessage(ValueError, msg):
+            City.country.get_prefetch_querysets(
+                instances=cities,
+                querysets=[Country.objects.all(), Country.objects.all()],
+            )
+
+    def test_get_prefetch_querysets_reverse_invalid_querysets_length(self):
+        usa = Country.objects.create(name="United States")
+        City.objects.create(name="Chicago")
+        countries = Country.objects.all()
+        msg = (
+            "querysets argument of get_prefetch_querysets() should have a length of 1."
+        )
+        with self.assertRaisesMessage(ValueError, msg):
+            usa.cities.get_prefetch_querysets(
+                instances=countries,
+                querysets=[City.objects.all(), City.objects.all()],
+            )
diff --git a/tests/one_to_one/tests.py b/tests/one_to_one/tests.py
index 65efee6074..83644871fe 100644
--- a/tests/one_to_one/tests.py
+++ b/tests/one_to_one/tests.py
@@ -1,5 +1,6 @@
 from django.db import IntegrityError, connection, transaction
 from django.test import TestCase
+from django.utils.deprecation import RemovedInDjango60Warning
 
 from .models import (
     Bar,
@@ -606,3 +607,23 @@ class OneToOneTests(TestCase):
         self.b1.place_id = self.p2.pk
         self.b1.save()
         self.assertEqual(self.b1.place, self.p2)
+
+    def test_get_prefetch_queryset_warning(self):
+        places = Place.objects.all()
+        msg = (
+            "get_prefetch_queryset() is deprecated. Use get_prefetch_querysets() "
+            "instead."
+        )
+        with self.assertWarnsMessage(RemovedInDjango60Warning, msg):
+            Place.bar.get_prefetch_queryset(places)
+
+    def test_get_prefetch_querysets_invalid_querysets_length(self):
+        places = Place.objects.all()
+        msg = (
+            "querysets argument of get_prefetch_querysets() should have a length of 1."
+        )
+        with self.assertRaisesMessage(ValueError, msg):
+            Place.bar.get_prefetch_querysets(
+                instances=places,
+                querysets=[Bar.objects.all(), Bar.objects.all()],
+            )
diff --git a/tests/prefetch_related/tests.py b/tests/prefetch_related/tests.py
index 4566de631e..dd56664f68 100644
--- a/tests/prefetch_related/tests.py
+++ b/tests/prefetch_related/tests.py
@@ -13,6 +13,7 @@ from django.test import (
     skipUnlessDBFeature,
 )
 from django.test.utils import CaptureQueriesContext
+from django.utils.deprecation import RemovedInDjango60Warning
 
 from .models import (
     Article,
@@ -1696,7 +1697,7 @@ class Ticket21760Tests(TestCase):
 
     def test_bug(self):
         prefetcher = get_prefetcher(self.rooms[0], "house", "house")[0]
-        queryset = prefetcher.get_prefetch_queryset(list(Room.objects.all()))[0]
+        queryset = prefetcher.get_prefetch_querysets(list(Room.objects.all()))[0]
         self.assertNotIn(" JOIN ", str(queryset.query))
 
 
@@ -1994,3 +1995,19 @@ class PrefetchLimitTests(TestDataMixin, TestCase):
         )
         with self.assertRaisesMessage(NotSupportedError, msg):
             list(Book.objects.prefetch_related(Prefetch("authors", authors[1:])))
+
+
+class GetCurrentQuerySetDeprecation(TestCase):
+    def test_get_current_queryset_warning(self):
+        msg = (
+            "Prefetch.get_current_queryset() is deprecated. Use "
+            "get_current_querysets() instead."
+        )
+        authors = Author.objects.all()
+        with self.assertWarnsMessage(RemovedInDjango60Warning, msg):
+            self.assertEqual(
+                Prefetch("authors", authors).get_current_queryset(1),
+                authors,
+            )
+        with self.assertWarnsMessage(RemovedInDjango60Warning, msg):
+            self.assertIsNone(Prefetch("authors").get_current_queryset(1))