Fix pk-only optimization for properties (#7142)

* Add callable/prop tests for pk-only optimization

* Fix related field pk-only optimization for props

Author: rpkilby

rest_framework/relations.py

@@ -175,8 +175,13 @@ def get_attribute(self, instance):
                 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

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)

tests/test_relations_pk.py

@@ -30,6 +30,25 @@ class Meta:
         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 Meta:
         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')