Custom CompositePrimaryKey and foreign key support on partitioned models

Photonios · Photonios · commit 552fb2712276 · 2025-07-10T21:23:24.000+02:00
diff --git a/psqlextra/models/partitioned.py b/psqlextra/models/partitioned.py
@@ -1,6 +1,10 @@
-from typing import Iterable
+from typing import Iterable, List, Optional, Tuple
 
+from django.core.exceptions import ImproperlyConfigured
+from django.db import models
 from django.db.models.base import ModelBase
+from django.db.models.fields.composite import CompositePrimaryKey
+from django.db.models.options import Options
 
 from psqlextra.types import PostgresPartitioningMethod
 
@@ -20,19 +24,140 @@ class PostgresPartitionedModelMeta(ModelBase):
     default_key: Iterable[str] = []
 
     def __new__(cls, name, bases, attrs, **kwargs):
-        new_class = super().__new__(cls, name, bases, attrs, **kwargs)
-        meta_class = attrs.pop("PartitioningMeta", None)
+        partitioning_meta_class = attrs.pop("PartitioningMeta", None)
+
+        partitioning_method = getattr(partitioning_meta_class, "method", None)
+        partitioning_key = getattr(partitioning_meta_class, "key", None)
+        special = getattr(partitioning_meta_class, "special", None)
 
-        method = getattr(meta_class, "method", None)
-        key = getattr(meta_class, "key", None)
+        if special:
+            cls._create_primary_key(attrs, partitioning_key)
 
         patitioning_meta = PostgresPartitionedModelOptions(
-            method=method or cls.default_method, key=key or cls.default_key
+            method=partitioning_method or cls.default_method,
+            key=partitioning_key or cls.default_key,
         )
 
+        new_class = super().__new__(cls, name, bases, attrs, **kwargs)
         new_class.add_to_class("_partitioning_meta", patitioning_meta)
         return new_class
 
+    @classmethod
+    def _create_primary_key(cls, attrs, partitioning_key: Optional[List[str]]):
+        pk = cls._find_primary_key(attrs)
+        if pk and isinstance(pk[1], CompositePrimaryKey):
+            return
+
+        if not pk:
+            attrs["id"] = attrs.get("id") or cls._create_auto_field(attrs)
+            pk_fields = ["id"]
+        else:
+            pk_fields = [pk[0]]
+
+        unique_pk_fields = set(pk_fields + (partitioning_key or []))
+        if len(unique_pk_fields) <= 1:
+            return
+
+        auto_generated_pk = CompositePrimaryKey(*sorted(unique_pk_fields))
+        attrs["pk"] = auto_generated_pk
+
+    @classmethod
+    def _create_auto_field(cls, attrs):
+        app_label = attrs.get("app_label")
+        meta_class = attrs.get("Meta", None)
+
+        pk_class = Options(meta_class, app_label)._get_default_pk_class()
+        return pk_class(verbose_name="ID", primary_key=True, auto_created=True)
+
+    @classmethod
+    def _find_primary_key(cls, attrs) -> Optional[Tuple[str, models.Field]]:
+        """Gets the field that has been marked by the user as the primary key
+        field for this model.
+
+        This is quite complex because Django allows a variety of options:
+
+        1. No PK at all. In this case, Django generates one named `id`
+           as an auto-increment integer (AutoField)
+
+        2. One field that has `primary_key=True`. Any field can have
+           this attribute, but Django would error if there were more.
+
+        3. One field named `pk`.
+
+        4. One field that has `primary_key=True` and a field that
+           is of type `CompositePrimaryKey` that includes that
+           field.
+
+        Since a table can only have one primary key, our goal here
+        is to find the field (if any) that is going to become
+        the primary key of the table.
+
+        Our logic is straight forward:
+
+        1. If there is a `CompositePrimaryKey`, that field becomes the primary key.
+
+        2. If there is a field with `primary_key=True`, that field becomes the primary key.
+
+        3. There is no primary key.
+        """
+
+        fields = {
+            name: value
+            for name, value in attrs.items()
+            if isinstance(value, models.Field)
+        }
+
+        fields_marked_as_pk = {
+            name: value for name, value in fields.items() if value.primary_key
+        }
+
+        # We cannot let the user define a field named `pk` that is not a CompositePrimaryKey
+        # already because when we generate a primary key, we want to name it `pk`.
+        field_named_pk = attrs.get("pk")
+        if field_named_pk and not field_named_pk.primary_key:
+            raise ImproperlyConfigured(
+                "You cannot define a field named `pk` that is not a primary key."
+            )
+
+        if field_named_pk:
+            if not isinstance(field_named_pk, CompositePrimaryKey):
+                raise ImproperlyConfigured(
+                    "You cannot define a field named `pk` that is not a composite primary key on a partitioned model. Either make `pk` a CompositePrimaryKey or rename it."
+                )
+
+            return ("pk", field_named_pk)
+
+        if not fields_marked_as_pk:
+            return None
+
+        # Make sure the user didn't define N primary keys. Django would also warn
+        # about this.
+        #
+        # One exception is a set up such as:
+        #
+        # >>> id = models.AutoField(primary_key=True)
+        # >>> timestamp = models.DateTimeField()
+        # >>> pk = models.CompositePrimaryKey("id", "timestamp")
+        #
+        # In this case, both `id` and `pk` are marked as primary key. Django
+        # allows this and just ignores the `primary_key=True` attribute
+        # on all the other fields except the composite one.
+        #
+        # We also handle this as expected and treat the CompositePrimaryKey
+        # as the primary key.
+        sorted_fields_marked_as_pk = sorted(
+            list(fields_marked_as_pk.items()),
+            key=lambda pair: 0
+            if isinstance(pair[1], CompositePrimaryKey)
+            else 1,
+        )
+        if len(sorted_fields_marked_as_pk[1:]) > 1:
+            raise ImproperlyConfigured(
+                "You cannot mark more than one fields as a primary key."
+            )
+
+        return sorted_fields_marked_as_pk[0]
+
 
 class PostgresPartitionedModel(
     PostgresModel, metaclass=PostgresPartitionedModelMeta
diff --git a/tests/test_partitioned_model.py b/tests/test_partitioned_model.py
@@ -1,7 +1,13 @@
+import django
+import pytest
+
+from django.core.exceptions import ImproperlyConfigured
+from django.db import models
+
 from psqlextra.models import PostgresPartitionedModel
 from psqlextra.types import PostgresPartitioningMethod
 
-from .fake_model import define_fake_partitioned_model
+from .fake_model import define_fake_model, define_fake_partitioned_model
 
 
 def test_partitioned_model_abstract():
@@ -70,3 +76,190 @@ def test_partitioned_model_key_option_none():
     model = define_fake_partitioned_model(partitioning_options=dict(key=None))
 
     assert model._partitioning_meta.key == []
+
+
+@pytest.mark.skipif(
+    django.VERSION < (5, 2),
+    reason="Django < 5.2 doesn't implement composite primary keys",
+)
+def test_partitioned_model_custom_composite_primary_key_with_auto_field():
+    model = define_fake_partitioned_model(
+        fields={
+            "auto_id": models.AutoField(),
+            "my_custom_pk": models.CompositePrimaryKey("auto_id", "timestamp"),
+            "timestamp": models.DateTimeField(),
+        },
+        partitioning_options=dict(key=["timestamp"], special=True),
+    )
+
+    assert isinstance(model._meta.pk, models.CompositePrimaryKey)
+    assert model._meta.pk.name == "my_custom_pk"
+    assert model._meta.pk.columns == ("auto_id", "timestamp")
+
+
+@pytest.mark.skipif(
+    django.VERSION < (5, 2),
+    reason="Django < 5.2 doesn't implement composite primary keys",
+)
+def test_partitioned_model_custom_composite_primary_key_with_id_field():
+    model = define_fake_partitioned_model(
+        fields={
+            "id": models.IntegerField(),
+            "my_custom_pk": models.CompositePrimaryKey("id", "timestamp"),
+            "timestamp": models.DateTimeField(),
+        },
+        partitioning_options=dict(key=["timestamp"], special=True),
+    )
+
+    assert isinstance(model._meta.pk, models.CompositePrimaryKey)
+    assert model._meta.pk.name == "my_custom_pk"
+    assert model._meta.pk.columns == ("id", "timestamp")
+
+
+@pytest.mark.skipif(
+    django.VERSION < (5, 2),
+    reason="Django < 5.2 doesn't implement composite primary keys",
+)
+def test_partitioned_model_custom_composite_primary_key_named_id():
+    model = define_fake_partitioned_model(
+        fields={
+            "other_field": models.TextField(),
+            "id": models.CompositePrimaryKey("other_field", "timestamp"),
+            "timestamp": models.DateTimeField(),
+        },
+        partitioning_options=dict(key=["timestamp"], special=True),
+    )
+
+    assert isinstance(model._meta.pk, models.CompositePrimaryKey)
+    assert model._meta.pk.name == "id"
+    assert model._meta.pk.columns == ("other_field", "timestamp")
+
+
+@pytest.mark.skipif(
+    django.VERSION < (5, 2),
+    reason="Django < 5.2 doesn't implement composite primary keys",
+)
+def test_partitioned_model_field_named_pk_not_composite_not_primary():
+    with pytest.raises(ImproperlyConfigured):
+        define_fake_partitioned_model(
+            fields={
+                "pk": models.TextField(),
+                "id": models.CompositePrimaryKey("other_field", "timestamp"),
+                "timestamp": models.DateTimeField(),
+            },
+            partitioning_options=dict(key=["timestamp"], special=True),
+        )
+
+
+@pytest.mark.skipif(
+    django.VERSION < (5, 2),
+    reason="Django < 5.2 doesn't implement composite primary keys",
+)
+def test_partitioned_model_field_named_pk_not_composite():
+    with pytest.raises(ImproperlyConfigured):
+        define_fake_partitioned_model(
+            fields={
+                "pk": models.AutoField(primary_key=True),
+                "timestamp": models.DateTimeField(),
+            },
+            partitioning_options=dict(key=["timestamp"], special=True),
+        )
+
+
+@pytest.mark.skipif(
+    django.VERSION < (5, 2),
+    reason="Django < 5.2 doesn't implement composite primary keys",
+)
+def test_partitioned_model_field_multiple_pks():
+    with pytest.raises(ImproperlyConfigured):
+        define_fake_partitioned_model(
+            fields={
+                "id": models.AutoField(primary_key=True),
+                "another_pk": models.TextField(primary_key=True),
+                "timestamp": models.DateTimeField(),
+                "real_pk": models.CompositePrimaryKey("id", "timestamp"),
+            },
+            partitioning_options=dict(key=["timestamp"], special=True),
+        )
+
+
+@pytest.mark.skipif(
+    django.VERSION < (5, 2),
+    reason="Django < 5.2 doesn't implement composite primary keys",
+)
+def test_partitioned_model_no_pk_defined():
+    model = define_fake_partitioned_model(
+        fields={
+            "timestamp": models.DateTimeField(),
+        },
+        partitioning_options=dict(key=["timestamp"], special=True),
+    )
+
+    assert isinstance(model._meta.pk, models.CompositePrimaryKey)
+    assert model._meta.pk.name == "pk"
+    assert model._meta.pk.columns == ("id", "timestamp")
+
+    id_field = model._meta.get_field("id")
+    assert id_field.name == "id"
+    assert id_field.column == "id"
+    assert isinstance(id_field, models.AutoField)
+    assert id_field.primary_key is True
+
+
+@pytest.mark.skipif(
+    django.VERSION < (5, 2),
+    reason="Django < 5.2 doesn't implement composite primary keys",
+)
+def test_partitioned_model_composite_primary_key():
+    model = define_fake_partitioned_model(
+        fields={
+            "id": models.AutoField(primary_key=True),
+            "pk": models.CompositePrimaryKey("id", "timestamp"),
+            "timestamp": models.DateTimeField(),
+        },
+        partitioning_options=dict(key=["timestamp"], special=True),
+    )
+
+    assert isinstance(model._meta.pk, models.CompositePrimaryKey)
+    assert model._meta.pk.name == "pk"
+    assert model._meta.pk.columns == ("id", "timestamp")
+
+
+@pytest.mark.skipif(
+    django.VERSION < (5, 2),
+    reason="Django < 5.2 doesn't implement composite primary keys",
+)
+def test_partitioned_model_composite_primary_key_foreign_key():
+    model = define_fake_partitioned_model(
+        fields={
+            "timestamp": models.DateTimeField(),
+        },
+        partitioning_options=dict(key=["timestamp"], special=True),
+    )
+
+    define_fake_model(
+        fields={
+            "model": models.ForeignKey(model, on_delete=models.CASCADE),
+        },
+    )
+
+
+@pytest.mark.skipif(
+    django.VERSION < (5, 2),
+    reason="Django < 5.2 doesn't implement composite primary keys",
+)
+def test_partitioned_model_custom_composite_primary_key_foreign_key():
+    model = define_fake_partitioned_model(
+        fields={
+            "id": models.TextField(primary_key=True),
+            "timestamp": models.DateTimeField(),
+            "custom": models.CompositePrimaryKey("id", "timestamp"),
+        },
+        partitioning_options=dict(key=["timestamp"], special=True),
+    )
+
+    define_fake_model(
+        fields={
+            "model": models.ForeignKey(model, on_delete=models.CASCADE),
+        },
+    )
