Fixed #30953 -- Made select_for_update() lock queryset's model when using "self" with multi-table inheritance.

Thanks Abhijeet Viswa for the report and initial patch.

2 files changed, 70 insertions, 1 deletions

diff --git a/tests/select_for_update/models.py b/tests/select_for_update/models.py
index b8154af3df..c84f9ad6b2 100644
--- a/tests/select_for_update/models.py
+++ b/tests/select_for_update/models.py
@@ -5,11 +5,20 @@ class Country(models.Model):
     name = models.CharField(max_length=30)
 
 
+class EUCountry(Country):
+    join_date = models.DateField()
+
+
 class City(models.Model):
     name = models.CharField(max_length=30)
     country = models.ForeignKey(Country, models.CASCADE)
 
 
+class EUCity(models.Model):
+    name = models.CharField(max_length=30)
+    country = models.ForeignKey(EUCountry, models.CASCADE)
+
+
 class Person(models.Model):
     name = models.CharField(max_length=30)
     born = models.ForeignKey(City, models.CASCADE, related_name='+')
diff --git a/tests/select_for_update/tests.py b/tests/select_for_update/tests.py
index 7859775cff..0bb21972d1 100644
--- a/tests/select_for_update/tests.py
+++ b/tests/select_for_update/tests.py
@@ -15,7 +15,7 @@ from django.test import (
 )
 from django.test.utils import CaptureQueriesContext
 
-from .models import City, Country, Person, PersonProfile
+from .models import City, Country, EUCity, EUCountry, Person, PersonProfile
 
 
 class SelectForUpdateTests(TransactionTestCase):
@@ -120,6 +120,47 @@ class SelectForUpdateTests(TransactionTestCase):
         self.assertTrue(self.has_for_update_sql(ctx.captured_queries, of=expected))
 
     @skipUnlessDBFeature('has_select_for_update_of')
+    def test_for_update_sql_model_inheritance_generated_of(self):
+        with transaction.atomic(), CaptureQueriesContext(connection) as ctx:
+            list(EUCountry.objects.select_for_update(of=('self',)))
+        if connection.features.select_for_update_of_column:
+            expected = ['select_for_update_eucountry"."country_ptr_id']
+        else:
+            expected = ['select_for_update_eucountry']
+        expected = [connection.ops.quote_name(value) for value in expected]
+        self.assertTrue(self.has_for_update_sql(ctx.captured_queries, of=expected))
+
+    @skipUnlessDBFeature('has_select_for_update_of')
+    def test_for_update_sql_model_inheritance_ptr_generated_of(self):
+        with transaction.atomic(), CaptureQueriesContext(connection) as ctx:
+            list(EUCountry.objects.select_for_update(of=('self', 'country_ptr',)))
+        if connection.features.select_for_update_of_column:
+            expected = [
+                'select_for_update_eucountry"."country_ptr_id',
+                'select_for_update_country"."id',
+            ]
+        else:
+            expected = ['select_for_update_eucountry', 'select_for_update_country']
+        expected = [connection.ops.quote_name(value) for value in expected]
+        self.assertTrue(self.has_for_update_sql(ctx.captured_queries, of=expected))
+
+    @skipUnlessDBFeature('has_select_for_update_of')
+    def test_for_update_sql_model_inheritance_nested_ptr_generated_of(self):
+        with transaction.atomic(), CaptureQueriesContext(connection) as ctx:
+            list(EUCity.objects.select_related('country').select_for_update(
+                of=('self', 'country__country_ptr',),
+            ))
+        if connection.features.select_for_update_of_column:
+            expected = [
+                'select_for_update_eucity"."id',
+                'select_for_update_country"."id',
+            ]
+        else:
+            expected = ['select_for_update_eucity', 'select_for_update_country']
+        expected = [connection.ops.quote_name(value) for value in expected]
+        self.assertTrue(self.has_for_update_sql(ctx.captured_queries, of=expected))
+
+    @skipUnlessDBFeature('has_select_for_update_of')
     def test_for_update_of_followed_by_values(self):
         with transaction.atomic():
             values = list(Person.objects.select_for_update(of=('self',)).values('pk'))
@@ -258,6 +299,25 @@ class SelectForUpdateTests(TransactionTestCase):
                         ).exclude(profile=None).select_for_update(of=(name,)).get()
 
     @skipUnlessDBFeature('has_select_for_update', 'has_select_for_update_of')
+    def test_model_inheritance_of_argument_raises_error_ptr_in_choices(self):
+        msg = (
+            'Invalid field name(s) given in select_for_update(of=(...)): '
+            'name. Only relational fields followed in the query are allowed. '
+            'Choices are: self, %s.'
+        )
+        with self.assertRaisesMessage(
+            FieldError,
+            msg % 'country, country__country_ptr',
+        ):
+            with transaction.atomic():
+                EUCity.objects.select_related(
+                    'country',
+                ).select_for_update(of=('name',)).get()
+        with self.assertRaisesMessage(FieldError, msg % 'country_ptr'):
+            with transaction.atomic():
+                EUCountry.objects.select_for_update(of=('name',)).get()
+
+    @skipUnlessDBFeature('has_select_for_update', 'has_select_for_update_of')
     def test_reverse_one_to_one_of_arguments(self):
         """
         Reverse OneToOneFields may be included in of=(...) as long as NULLs