Fixed #373 -- Added CompositePrimaryKey.

Thanks Lily Foote and Simon Charette for reviews and mentoring
this Google Summer of Code 2024 project.

Co-authored-by: Simon Charette <charette.s@gmail.com>
Co-authored-by: Lily Foote <code@lilyf.org>

4 files changed, 185 insertions, 0 deletions

diff --git a/tests/migrations/test_autodetector.py b/tests/migrations/test_autodetector.py
index de62170eb3..33196ea6f4 100644
--- a/tests/migrations/test_autodetector.py
+++ b/tests/migrations/test_autodetector.py
@@ -5059,6 +5059,95 @@ class AutodetectorTests(BaseAutodetectorTests):
         self.assertOperationTypes(changes, "testapp", 0, ["CreateModel"])
         self.assertOperationAttributes(changes, "testapp", 0, 0, name="Book")
 
+    @mock.patch(
+        "django.db.migrations.questioner.MigrationQuestioner.ask_not_null_addition"
+    )
+    def test_add_composite_pk(self, mocked_ask_method):
+        before = [
+            ModelState(
+                "app",
+                "foo",
+                [
+                    ("id", models.AutoField(primary_key=True)),
+                ],
+            ),
+        ]
+        after = [
+            ModelState(
+                "app",
+                "foo",
+                [
+                    ("pk", models.CompositePrimaryKey("foo_id", "bar_id")),
+                    ("id", models.IntegerField()),
+                ],
+            ),
+        ]
+
+        changes = self.get_changes(before, after)
+        self.assertEqual(mocked_ask_method.call_count, 0)
+        self.assertNumberMigrations(changes, "app", 1)
+        self.assertOperationTypes(changes, "app", 0, ["AddField", "AlterField"])
+        self.assertOperationAttributes(
+            changes,
+            "app",
+            0,
+            0,
+            name="pk",
+            model_name="foo",
+            preserve_default=True,
+        )
+        self.assertOperationAttributes(
+            changes,
+            "app",
+            0,
+            1,
+            name="id",
+            model_name="foo",
+            preserve_default=True,
+        )
+
+    def test_remove_composite_pk(self):
+        before = [
+            ModelState(
+                "app",
+                "foo",
+                [
+                    ("pk", models.CompositePrimaryKey("foo_id", "bar_id")),
+                    ("id", models.IntegerField()),
+                ],
+            ),
+        ]
+        after = [
+            ModelState(
+                "app",
+                "foo",
+                [
+                    ("id", models.AutoField(primary_key=True)),
+                ],
+            ),
+        ]
+
+        changes = self.get_changes(before, after)
+        self.assertNumberMigrations(changes, "app", 1)
+        self.assertOperationTypes(changes, "app", 0, ["RemoveField", "AlterField"])
+        self.assertOperationAttributes(
+            changes,
+            "app",
+            0,
+            0,
+            name="pk",
+            model_name="foo",
+        )
+        self.assertOperationAttributes(
+            changes,
+            "app",
+            0,
+            1,
+            name="id",
+            model_name="foo",
+            preserve_default=True,
+        )
+
 
 class MigrationSuggestNameTests(SimpleTestCase):
     def test_no_operations(self):
diff --git a/tests/migrations/test_operations.py b/tests/migrations/test_operations.py
index da0ec93dcd..6312a7d4a2 100644
--- a/tests/migrations/test_operations.py
+++ b/tests/migrations/test_operations.py
@@ -6287,6 +6287,61 @@ class OperationTests(OperationTestBase):
         self.assertEqual(pony_new.generated, 1)
         self.assertEqual(pony_new.static, 2)
 
+    def test_composite_pk_operations(self):
+        app_label = "test_d8d90af6"
+        project_state = self.set_up_test_model(app_label)
+        operation_1 = migrations.AddField(
+            "Pony", "pk", models.CompositePrimaryKey("id", "pink")
+        )
+        operation_2 = migrations.AlterField("Pony", "id", models.IntegerField())
+        operation_3 = migrations.RemoveField("Pony", "pk")
+        table_name = f"{app_label}_pony"
+
+        # 1. Add field (pk).
+        new_state = project_state.clone()
+        operation_1.state_forwards(app_label, new_state)
+        with connection.schema_editor() as editor:
+            operation_1.database_forwards(app_label, editor, project_state, new_state)
+        self.assertColumnNotExists(table_name, "pk")
+        Pony = new_state.apps.get_model(app_label, "pony")
+        obj_1 = Pony.objects.create(weight=1)
+        msg = (
+            f"obj_1={obj_1}, "
+            f"obj_1.id={obj_1.id}, "
+            f"obj_1.pink={obj_1.pink}, "
+            f"obj_1.pk={obj_1.pk}, "
+            f"Pony._meta.pk={repr(Pony._meta.pk)}, "
+            f"Pony._meta.get_field('id')={repr(Pony._meta.get_field('id'))}"
+        )
+        self.assertEqual(obj_1.pink, 3, msg)
+        self.assertEqual(obj_1.pk, (obj_1.id, obj_1.pink), msg)
+
+        # 2. Alter field (id -> IntegerField()).
+        project_state, new_state = new_state, new_state.clone()
+        operation_2.state_forwards(app_label, new_state)
+        with connection.schema_editor() as editor:
+            operation_2.database_forwards(app_label, editor, project_state, new_state)
+        Pony = new_state.apps.get_model(app_label, "pony")
+        obj_1 = Pony.objects.get(id=obj_1.id)
+        self.assertEqual(obj_1.pink, 3)
+        self.assertEqual(obj_1.pk, (obj_1.id, obj_1.pink))
+        obj_2 = Pony.objects.create(id=2, weight=2)
+        self.assertEqual(obj_2.id, 2)
+        self.assertEqual(obj_2.pink, 3)
+        self.assertEqual(obj_2.pk, (obj_2.id, obj_2.pink))
+
+        # 3. Remove field (pk).
+        project_state, new_state = new_state, new_state.clone()
+        operation_3.state_forwards(app_label, new_state)
+        with connection.schema_editor() as editor:
+            operation_3.database_forwards(app_label, editor, project_state, new_state)
+        Pony = new_state.apps.get_model(app_label, "pony")
+        obj_1 = Pony.objects.get(id=obj_1.id)
+        self.assertEqual(obj_1.pk, obj_1.id)
+        obj_2 = Pony.objects.get(id=obj_2.id)
+        self.assertEqual(obj_2.id, 2)
+        self.assertEqual(obj_2.pk, obj_2.id)
+
 
 class SwappableOperationTests(OperationTestBase):
     """
diff --git a/tests/migrations/test_state.py b/tests/migrations/test_state.py
index dbbdf77734..d6ecaa1c5d 100644
--- a/tests/migrations/test_state.py
+++ b/tests/migrations/test_state.py
@@ -1206,6 +1206,28 @@ class StateTests(SimpleTestCase):
         choices_field = Author._meta.get_field("choice")
         self.assertEqual(list(choices_field.choices), choices)
 
+    def test_composite_pk_state(self):
+        new_apps = Apps(["migrations"])
+
+        class Foo(models.Model):
+            pk = models.CompositePrimaryKey("account_id", "id")
+            account_id = models.SmallIntegerField()
+            id = models.SmallIntegerField()
+
+            class Meta:
+                app_label = "migrations"
+                apps = new_apps
+
+        project_state = ProjectState.from_apps(new_apps)
+        model_state = project_state.models["migrations", "foo"]
+        self.assertEqual(len(model_state.options), 2)
+        self.assertEqual(model_state.options["constraints"], [])
+        self.assertEqual(model_state.options["indexes"], [])
+        self.assertEqual(len(model_state.fields), 3)
+        self.assertIn("pk", model_state.fields)
+        self.assertIn("account_id", model_state.fields)
+        self.assertIn("id", model_state.fields)
+
 
 class StateRelationsTests(SimpleTestCase):
     def get_base_project_state(self):
diff --git a/tests/migrations/test_writer.py b/tests/migrations/test_writer.py
index 51783b7346..953a3cdb6c 100644
--- a/tests/migrations/test_writer.py
+++ b/tests/migrations/test_writer.py
@@ -1138,3 +1138,22 @@ class WriterTests(SimpleTestCase):
             ValueError, "'TestModel1' must inherit from 'BaseSerializer'."
         ):
             MigrationWriter.register_serializer(complex, TestModel1)
+
+    def test_composite_pk_import(self):
+        migration = type(
+            "Migration",
+            (migrations.Migration,),
+            {
+                "operations": [
+                    migrations.AddField(
+                        "foo",
+                        "bar",
+                        models.CompositePrimaryKey("foo_id", "bar_id"),
+                    ),
+                ],
+            },
+        )
+        writer = MigrationWriter(migration)
+        output = writer.as_string()
+        self.assertEqual(output.count("import"), 1)
+        self.assertIn("from django.db import migrations, models", output)