Fixed #31246 -- Fixed locking models in QuerySet.select_for_update(of=()) for related fields and parent link fields with multi-table inheritance.

Partly regression in 0107e3d1058f653f66032f7fd3a0bd61e96bf782.

2 files changed, 47 insertions, 7 deletions

diff --git a/tests/select_for_update/models.py b/tests/select_for_update/models.py
index c84f9ad6b2..305e8cac49 100644
--- a/tests/select_for_update/models.py
+++ b/tests/select_for_update/models.py
@@ -1,7 +1,11 @@
 from django.db import models
 
 
-class Country(models.Model):
+class Entity(models.Model):
+    pass
+
+
+class Country(Entity):
     name = models.CharField(max_length=30)
 
 
diff --git a/tests/select_for_update/tests.py b/tests/select_for_update/tests.py
index 0bb21972d1..3622a95c11 100644
--- a/tests/select_for_update/tests.py
+++ b/tests/select_for_update/tests.py
@@ -113,7 +113,10 @@ class SelectForUpdateTests(TransactionTestCase):
             ))
         features = connections['default'].features
         if features.select_for_update_of_column:
-            expected = ['select_for_update_person"."id', 'select_for_update_country"."id']
+            expected = [
+                'select_for_update_person"."id',
+                'select_for_update_country"."entity_ptr_id',
+            ]
         else:
             expected = ['select_for_update_person', 'select_for_update_country']
         expected = [connection.ops.quote_name(value) for value in expected]
@@ -137,7 +140,7 @@ class SelectForUpdateTests(TransactionTestCase):
         if connection.features.select_for_update_of_column:
             expected = [
                 'select_for_update_eucountry"."country_ptr_id',
-                'select_for_update_country"."id',
+                'select_for_update_country"."entity_ptr_id',
             ]
         else:
             expected = ['select_for_update_eucountry', 'select_for_update_country']
@@ -145,6 +148,22 @@ 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_related_model_inheritance_generated_of(self):
+        with transaction.atomic(), CaptureQueriesContext(connection) as ctx:
+            list(EUCity.objects.select_related('country').select_for_update(
+                of=('self', 'country'),
+            ))
+        if connection.features.select_for_update_of_column:
+            expected = [
+                'select_for_update_eucity"."id',
+                'select_for_update_eucountry"."country_ptr_id',
+            ]
+        else:
+            expected = ['select_for_update_eucity', '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_nested_ptr_generated_of(self):
         with transaction.atomic(), CaptureQueriesContext(connection) as ctx:
             list(EUCity.objects.select_related('country').select_for_update(
@@ -153,7 +172,7 @@ class SelectForUpdateTests(TransactionTestCase):
         if connection.features.select_for_update_of_column:
             expected = [
                 'select_for_update_eucity"."id',
-                'select_for_update_country"."id',
+                'select_for_update_country"."entity_ptr_id',
             ]
         else:
             expected = ['select_for_update_eucity', 'select_for_update_country']
@@ -161,6 +180,22 @@ 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_multilevel_model_inheritance_ptr_generated_of(self):
+        with transaction.atomic(), CaptureQueriesContext(connection) as ctx:
+            list(EUCountry.objects.select_for_update(
+                of=('country_ptr', 'country_ptr__entity_ptr'),
+            ))
+        if connection.features.select_for_update_of_column:
+            expected = [
+                'select_for_update_country"."entity_ptr_id',
+                'select_for_update_entity"."id',
+            ]
+        else:
+            expected = ['select_for_update_country', 'select_for_update_entity']
+        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'))
@@ -264,7 +299,8 @@ class SelectForUpdateTests(TransactionTestCase):
         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.'
+            'Choices are: self, born, born__country, '
+            'born__country__entity_ptr.'
         )
         invalid_of = [
             ('nonexistent',),
@@ -307,13 +343,13 @@ class SelectForUpdateTests(TransactionTestCase):
         )
         with self.assertRaisesMessage(
             FieldError,
-            msg % 'country, country__country_ptr',
+            msg % 'country, country__country_ptr, country__country_ptr__entity_ptr',
         ):
             with transaction.atomic():
                 EUCity.objects.select_related(
                     'country',
                 ).select_for_update(of=('name',)).get()
-        with self.assertRaisesMessage(FieldError, msg % 'country_ptr'):
+        with self.assertRaisesMessage(FieldError, msg % 'country_ptr, country_ptr__entity_ptr'):
             with transaction.atomic():
                 EUCountry.objects.select_for_update(of=('name',)).get()