From 22ce5d0031bd795ade081394043833e82046016c Mon Sep 17 00:00:00 2001
From: Mariusz Felisiak <felisiak.mariusz@gmail.com>
Date: Mon, 23 Dec 2019 23:28:59 +0100
Subject: [PATCH] Fixed #31106 -- Fixed migrations crash on PostgreSQL 10+ when
 adding FK constraints inline and changing data.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This allows adding foreign key constraints inline and changing data in
the same migration on PostgreSQL 10+.

Regression in 738faf9da2a5cd03148a36375db80746c99c9623.

Thanks Janne Rönkkö for the report and Simon Charette for the
implementation idea and review.
---
 django/db/backends/postgresql/schema.py |  7 ++++-
 docs/releases/3.0.2.txt                 |  4 +++
 tests/schema/tests.py                   | 37 +++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/django/db/backends/postgresql/schema.py b/django/db/backends/postgresql/schema.py
index cf90bb8f81..e5a0008025 100644
--- a/django/db/backends/postgresql/schema.py
+++ b/django/db/backends/postgresql/schema.py
@@ -19,7 +19,12 @@ class DatabaseSchemaEditor(BaseDatabaseSchemaEditor):
     sql_delete_index = "DROP INDEX IF EXISTS %(name)s"
     sql_delete_index_concurrently = "DROP INDEX CONCURRENTLY IF EXISTS %(name)s"
 
-    sql_create_column_inline_fk = 'REFERENCES %(to_table)s(%(to_column)s)%(deferrable)s'
+    # Setting the constraint to IMMEDIATE to allow changing data in the same
+    # transaction.
+    sql_create_column_inline_fk = (
+        'CONSTRAINT %(name)s REFERENCES %(to_table)s(%(to_column)s)%(deferrable)s'
+        '; SET CONSTRAINTS %(name)s IMMEDIATE'
+    )
     # Setting the constraint to IMMEDIATE runs any deferred checks to allow
     # dropping it in the same transaction.
     sql_delete_fk = "SET CONSTRAINTS %(name)s IMMEDIATE; ALTER TABLE %(table)s DROP CONSTRAINT %(name)s"
diff --git a/docs/releases/3.0.2.txt b/docs/releases/3.0.2.txt
index ccc82b18ae..ebb85d1f46 100644
--- a/docs/releases/3.0.2.txt
+++ b/docs/releases/3.0.2.txt
@@ -14,3 +14,7 @@ Bugfixes
 
 * Fixed a regression in Django 3.0 where ``QuerySet.exists()`` crashed if a
   queryset contained an aggregation over a ``Subquery()`` (:ticket:`31109`).
+
+* Fixed a regression in Django 3.0 that caused a migration crash on PostgreSQL
+  10+ when adding a foreign key and changing data in the same migration
+  (:ticket:`31106`).
diff --git a/tests/schema/tests.py b/tests/schema/tests.py
index b4d95c431c..d8f23606c1 100644
--- a/tests/schema/tests.py
+++ b/tests/schema/tests.py
@@ -275,6 +275,43 @@ class SchemaTests(TransactionTestCase):
             if sql.startswith('ALTER TABLE') and 'ADD CONSTRAINT' in sql
         ])
 
+    @skipUnlessDBFeature('can_create_inline_fk')
+    def test_add_inline_fk_update_data(self):
+        with connection.schema_editor() as editor:
+            editor.create_model(Node)
+        # Add an inline foreign key and update data in the same transaction.
+        new_field = ForeignKey(Node, CASCADE, related_name='new_fk', null=True)
+        new_field.set_attributes_from_name('new_parent_fk')
+        parent = Node.objects.create()
+        with connection.schema_editor() as editor:
+            editor.add_field(Node, new_field)
+            editor.execute('UPDATE schema_node SET new_parent_fk_id = %s;', [parent.pk])
+        self.assertIn('new_parent_fk_id', self.get_indexes(Node._meta.db_table))
+
+    @skipUnlessDBFeature(
+        'can_create_inline_fk',
+        'allows_multiple_constraints_on_same_fields',
+    )
+    @isolate_apps('schema')
+    def test_add_inline_fk_index_update_data(self):
+        class Node(Model):
+            class Meta:
+                app_label = 'schema'
+
+        with connection.schema_editor() as editor:
+            editor.create_model(Node)
+        # Add an inline foreign key, update data, and an index in the same
+        # transaction.
+        new_field = ForeignKey(Node, CASCADE, related_name='new_fk', null=True)
+        new_field.set_attributes_from_name('new_parent_fk')
+        parent = Node.objects.create()
+        with connection.schema_editor() as editor:
+            editor.add_field(Node, new_field)
+            Node._meta.add_field(new_field)
+            editor.execute('UPDATE schema_node SET new_parent_fk_id = %s;', [parent.pk])
+            editor.add_index(Node, Index(fields=['new_parent_fk'], name='new_parent_inline_fk_idx'))
+        self.assertIn('new_parent_fk_id', self.get_indexes(Node._meta.db_table))
+
     @skipUnlessDBFeature('supports_foreign_keys')
     def test_char_field_with_db_index_to_fk(self):
         # Create the table