Skip to content

Commit

Permalink
Use attname to get pk value (#650)
Browse files Browse the repository at this point in the history
* Use `attname` to get pk value

* Add tests for model as primary key

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add clarifying comments to _get_pk_value

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
evanandrews-xrd and pre-commit-ci[bot] committed Jun 11, 2024
1 parent 92805e2 commit 2c0bd0f
Show file tree
Hide file tree
Showing 3 changed files with 63 additions and 5 deletions.
7 changes: 4 additions & 3 deletions auditlog/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -213,12 +213,13 @@ def _get_pk_value(self, instance):
:type instance: Model
:return: The primary key value of the given model instance.
"""
pk_field = instance._meta.pk.name
# Should be equivalent to `instance.pk`.
pk_field = instance._meta.pk.attname
pk = getattr(instance, pk_field, None)

# Check to make sure that we got a pk not a model object.
if isinstance(pk, models.Model):
pk = self._get_pk_value(pk)
# Should be guaranteed as we used `attname` above, not `name`.
assert not isinstance(pk, models.Model)
return pk

def _get_serialized_data_or_none(self, instance):
Expand Down
21 changes: 21 additions & 0 deletions auditlog_tests/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,26 @@ class UUIDPrimaryKeyModel(models.Model):
history = AuditlogHistoryField(delete_related=True, pk_indexable=False)


class ModelPrimaryKeyModel(models.Model):
"""
A model with another model as primary key.
"""

key = models.OneToOneField(
"SimpleModel",
primary_key=True,
on_delete=models.CASCADE,
related_name="reverse_primary_key",
)

text = models.TextField(blank=True)
boolean = models.BooleanField(default=False)
integer = models.IntegerField(blank=True, null=True)
datetime = models.DateTimeField(auto_now=True)

history = AuditlogHistoryField(delete_related=True, pk_indexable=False)


class ProxyModel(SimpleModel):
"""
A model that is a proxy for another model.
Expand Down Expand Up @@ -338,6 +358,7 @@ class AutoManyRelatedModel(models.Model):

auditlog.register(AltPrimaryKeyModel)
auditlog.register(UUIDPrimaryKeyModel)
auditlog.register(ModelPrimaryKeyModel)
auditlog.register(ProxyModel)
auditlog.register(RelatedModel)
auditlog.register(ManyRelatedModel)
Expand Down
40 changes: 38 additions & 2 deletions auditlog_tests/tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
from django.db.models import JSONField, Value
from django.db.models.functions import Now
from django.db.models.signals import pre_save
from django.test import RequestFactory, TestCase, override_settings
from django.test import RequestFactory, TestCase, TransactionTestCase, override_settings
from django.urls import resolve, reverse
from django.utils import dateformat, formats
from django.utils import timezone as django_timezone
Expand All @@ -46,6 +46,7 @@
JSONModel,
ManyRelatedModel,
ManyRelatedOtherModel,
ModelPrimaryKeyModel,
NoDeleteHistoryModel,
NullableJSONModel,
PostgresArrayFieldModel,
Expand Down Expand Up @@ -332,6 +333,41 @@ class UUIDPrimaryKeyModelModelWithActorTest(
pass


class ModelPrimaryKeyModelBase(SimpleModelTest):
def make_object(self):
self.key = super().make_object()
return ModelPrimaryKeyModel.objects.create(key=self.key, text="I am strange.")


class ModelPrimaryKeyModelTest(NoActorMixin, ModelPrimaryKeyModelBase):
pass


class ModelPrimaryKeyModelWithActorTest(WithActorMixin, ModelPrimaryKeyModelBase):
pass


# Must inherit from TransactionTestCase to use self.assertNumQueries.
class ModelPrimaryKeyTest(TransactionTestCase):
def test_get_pk_value(self):
"""
Test that the primary key can be retrieved without additional database queries.
"""
key = SimpleModel.objects.create(text="I am not difficult.")
obj = ModelPrimaryKeyModel.objects.create(key=key, text="I am strange.")
# Refresh the object so the primary key object is not cached.
obj.refresh_from_db()
with self.assertNumQueries(0):
pk = LogEntry.objects._get_pk_value(obj)
self.assertEqual(pk, obj.pk)
self.assertEqual(pk, key.pk)
# Sanity check: verify accessing obj.key causes database access.
with self.assertNumQueries(1):
pk = obj.key.pk
self.assertEqual(pk, obj.pk)
self.assertEqual(pk, key.pk)


class ProxyModelBase(SimpleModelTest):
def make_object(self):
return ProxyModel.objects.create(text="I am not what you think.")
Expand Down Expand Up @@ -1206,7 +1242,7 @@ def test_register_models_register_app(self):

self.assertTrue(self.test_auditlog.contains(SimpleExcludeModel))
self.assertTrue(self.test_auditlog.contains(ChoicesFieldModel))
self.assertEqual(len(self.test_auditlog.get_models()), 26)
self.assertEqual(len(self.test_auditlog.get_models()), 27)

def test_register_models_register_model_with_attrs(self):
self.test_auditlog._register_models(
Expand Down

0 comments on commit 2c0bd0f

Please sign in to comment.