From f323049ecc24b48b8453867d0d8952e8e0c1c003 Mon Sep 17 00:00:00 2001
From: Ryan P Kilby <kilbyr@gmail.com>
Date: Thu, 3 Sep 2020 03:49:15 -0700
Subject: [PATCH] Fix pk-only optimization for properties (#7142)

* Add callable/prop tests for pk-only optimization

* Fix related field pk-only optimization for props
---
 rest_framework/relations.py |  9 +++++--
 tests/models.py             |  9 +++++++
 tests/test_relations_pk.py  | 47 +++++++++++++++++++++++++++++++++++++
 3 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/rest_framework/relations.py b/rest_framework/relations.py
index 3a2a8fb4b..3cd46379d 100644
--- a/rest_framework/relations.py
+++ b/rest_framework/relations.py
@@ -175,8 +175,13 @@ class RelatedField(Field):
                 value = attribute_instance.serializable_value(self.source_attrs[-1])
                 if is_simple_callable(value):
                     # Handle edge case where the relationship `source` argument
-                    # points to a `get_relationship()` method on the model
-                    value = value().pk
+                    # points to a `get_relationship()` method on the model.
+                    value = value()
+
+                # Handle edge case where relationship `source` argument points
+                # to an instance instead of a pk (e.g., a `@property`).
+                value = getattr(value, 'pk', value)
+
                 return PKOnlyObject(pk=value)
             except AttributeError:
                 pass
diff --git a/tests/models.py b/tests/models.py
index 2c340a3e6..afe649760 100644
--- a/tests/models.py
+++ b/tests/models.py
@@ -37,6 +37,15 @@ class ManyToManySource(RESTFrameworkModel):
 class ForeignKeyTarget(RESTFrameworkModel):
     name = models.CharField(max_length=100)
 
+    def get_first_source(self):
+        """Used for testing related field against a callable."""
+        return self.sources.all().order_by('pk')[0]
+
+    @property
+    def first_source(self):
+        """Used for testing related field against a property."""
+        return self.sources.all().order_by('pk')[0]
+
 
 class UUIDForeignKeyTarget(RESTFrameworkModel):
     uuid = models.UUIDField(primary_key=True, default=uuid.uuid4)
diff --git a/tests/test_relations_pk.py b/tests/test_relations_pk.py
index 0da9da890..7a4878a2b 100644
--- a/tests/test_relations_pk.py
+++ b/tests/test_relations_pk.py
@@ -30,6 +30,25 @@ class ForeignKeyTargetSerializer(serializers.ModelSerializer):
         fields = ('id', 'name', 'sources')
 
 
+class ForeignKeyTargetCallableSourceSerializer(serializers.ModelSerializer):
+    first_source = serializers.PrimaryKeyRelatedField(
+        source='get_first_source',
+        read_only=True,
+    )
+
+    class Meta:
+        model = ForeignKeyTarget
+        fields = ('id', 'name', 'first_source')
+
+
+class ForeignKeyTargetPropertySourceSerializer(serializers.ModelSerializer):
+    first_source = serializers.PrimaryKeyRelatedField(read_only=True)
+
+    class Meta:
+        model = ForeignKeyTarget
+        fields = ('id', 'name', 'first_source')
+
+
 class ForeignKeySourceSerializer(serializers.ModelSerializer):
     class Meta:
         model = ForeignKeySource
@@ -389,6 +408,34 @@ class PKForeignKeyTests(TestCase):
         assert len(queryset) == 1
 
 
+class PKRelationTests(TestCase):
+
+    def setUp(self):
+        self.target = ForeignKeyTarget.objects.create(name='target-1')
+        ForeignKeySource.objects.create(name='source-1', target=self.target)
+        ForeignKeySource.objects.create(name='source-2', target=self.target)
+
+    def test_relation_field_callable_source(self):
+        serializer = ForeignKeyTargetCallableSourceSerializer(self.target)
+        expected = {
+            'id': 1,
+            'name': 'target-1',
+            'first_source': 1,
+        }
+        with self.assertNumQueries(1):
+            self.assertEqual(serializer.data, expected)
+
+    def test_relation_field_property_source(self):
+        serializer = ForeignKeyTargetPropertySourceSerializer(self.target)
+        expected = {
+            'id': 1,
+            'name': 'target-1',
+            'first_source': 1,
+        }
+        with self.assertNumQueries(1):
+            self.assertEqual(serializer.data, expected)
+
+
 class PKNullableForeignKeyTests(TestCase):
     def setUp(self):
         target = ForeignKeyTarget(name='target-1')