Fixed #28010 -- Added FOR UPDATE OF support to QuerySet.select_for_update().

2 files changed, 94 insertions, 3 deletions

diff --git a/tests/select_for_update/models.py b/tests/select_for_update/models.py
index 48ad58faa9..b04ed31b00 100644
--- a/tests/select_for_update/models.py
+++ b/tests/select_for_update/models.py
@@ -1,5 +1,16 @@
 from django.db import models
 
 
+class Country(models.Model):
+    name = models.CharField(max_length=30)
+
+
+class City(models.Model):
+    name = models.CharField(max_length=30)
+    country = models.ForeignKey(Country, models.CASCADE)
+
+
 class Person(models.Model):
     name = models.CharField(max_length=30)
+    born = models.ForeignKey(City, models.CASCADE, related_name='+')
+    died = models.ForeignKey(City, models.CASCADE, related_name='+')
diff --git a/tests/select_for_update/tests.py b/tests/select_for_update/tests.py
index 0c581f0f37..7228af6e8e 100644
--- a/tests/select_for_update/tests.py
+++ b/tests/select_for_update/tests.py
@@ -4,6 +4,7 @@ from unittest import mock
 
 from multiple_database.routers import TestRouter
 
+from django.core.exceptions import FieldError
 from django.db import (
     DatabaseError, NotSupportedError, connection, connections, router,
     transaction,
@@ -14,7 +15,7 @@ from django.test import (
 )
 from django.test.utils import CaptureQueriesContext
 
-from .models import Person
+from .models import City, Country, Person
 
 
 class SelectForUpdateTests(TransactionTestCase):
@@ -24,7 +25,11 @@ class SelectForUpdateTests(TransactionTestCase):
     def setUp(self):
         # This is executed in autocommit mode so that code in
         # run_select_for_update can see this data.
-        self.person = Person.objects.create(name='Reinhardt')
+        self.country1 = Country.objects.create(name='Belgium')
+        self.country2 = Country.objects.create(name='France')
+        self.city1 = City.objects.create(name='Liberchies', country=self.country1)
+        self.city2 = City.objects.create(name='Samois-sur-Seine', country=self.country2)
+        self.person = Person.objects.create(name='Reinhardt', born=self.city1, died=self.city2)
 
         # We need another database connection in transaction to test that one
         # connection issuing a SELECT ... FOR UPDATE will block.
@@ -90,6 +95,29 @@ class SelectForUpdateTests(TransactionTestCase):
             list(Person.objects.all().select_for_update(skip_locked=True))
         self.assertTrue(self.has_for_update_sql(ctx.captured_queries, skip_locked=True))
 
+    @skipUnlessDBFeature('has_select_for_update_of')
+    def test_for_update_sql_generated_of(self):
+        """
+        The backend's FOR UPDATE OF variant appears in the generated SQL when
+        select_for_update() is invoked.
+        """
+        with transaction.atomic(), CaptureQueriesContext(connection) as ctx:
+            list(Person.objects.select_related(
+                'born__country',
+            ).select_for_update(
+                of=('born__country',),
+            ).select_for_update(
+                of=('self', 'born__country')
+            ))
+        features = connections['default'].features
+        if features.select_for_update_of_column:
+            expected = ['"select_for_update_person"."id"', '"select_for_update_country"."id"']
+        else:
+            expected = ['"select_for_update_person"', '"select_for_update_country"']
+        if features.uppercases_column_names:
+            expected = [value.upper() for value in expected]
+        self.assertTrue(self.has_for_update_sql(ctx.captured_queries, of=expected))
+
     @skipUnlessDBFeature('has_select_for_update_nowait')
     def test_nowait_raises_error_on_block(self):
         """
@@ -152,6 +180,58 @@ class SelectForUpdateTests(TransactionTestCase):
             with transaction.atomic():
                 Person.objects.select_for_update(skip_locked=True).get()
 
+    @skipIfDBFeature('has_select_for_update_of')
+    @skipUnlessDBFeature('has_select_for_update')
+    def test_unsupported_of_raises_error(self):
+        """
+        NotSupportedError is raised if a SELECT...FOR UPDATE OF... is run on
+        a database backend that supports FOR UPDATE but not OF.
+        """
+        msg = 'FOR UPDATE OF is not supported on this database backend.'
+        with self.assertRaisesMessage(NotSupportedError, msg):
+            with transaction.atomic():
+                Person.objects.select_for_update(of=('self',)).get()
+
+    @skipUnlessDBFeature('has_select_for_update', 'has_select_for_update_of')
+    def test_unrelated_of_argument_raises_error(self):
+        """
+        FieldError is raised if a non-relation field is specified in of=(...).
+        """
+        msg = (
+            'Invalid field name(s) given in select_for_update(of=(...)): %s. '
+            'Only relational fields followed in the query are allowed. '
+            'Choices are: self, born, born__country.'
+        )
+        invalid_of = [
+            ('nonexistent',),
+            ('name',),
+            ('born__nonexistent',),
+            ('born__name',),
+            ('born__nonexistent', 'born__name'),
+        ]
+        for of in invalid_of:
+            with self.subTest(of=of):
+                with self.assertRaisesMessage(FieldError, msg % ', '.join(of)):
+                    with transaction.atomic():
+                        Person.objects.select_related('born__country').select_for_update(of=of).get()
+
+    @skipUnlessDBFeature('has_select_for_update', 'has_select_for_update_of')
+    def test_related_but_unselected_of_argument_raises_error(self):
+        """
+        FieldError is raised if a relation field that is not followed in the
+        query is specified in of=(...).
+        """
+        msg = (
+            'Invalid field name(s) given in select_for_update(of=(...)): %s. '
+            'Only relational fields followed in the query are allowed. '
+            'Choices are: self, born.'
+        )
+        for name in ['born__country', 'died', 'died__country']:
+            with self.subTest(name=name):
+                with self.assertRaisesMessage(FieldError, msg % name):
+                    with transaction.atomic():
+                        Person.objects.select_related('born').select_for_update(of=(name,)).get()
+
     @skipUnlessDBFeature('has_select_for_update')
     def test_for_update_after_from(self):
         features_class = connections['default'].features.__class__
@@ -182,7 +262,7 @@ class SelectForUpdateTests(TransactionTestCase):
 
     @skipUnlessDBFeature('supports_select_for_update_with_limit')
     def test_select_for_update_with_limit(self):
-        other = Person.objects.create(name='Grappeli')
+        other = Person.objects.create(name='Grappeli', born=self.city1, died=self.city2)
         with transaction.atomic():
             qs = list(Person.objects.all().order_by('pk').select_for_update()[1:2])
             self.assertEqual(qs[0], other)