Merge pull request #57 from PeterJCLaw/fix-non-id-primary-keys

Have `PostgresQuerySet.insert` always return the primary key

Author: Photonios

psqlextra/manager/manager.py

@@ -187,12 +187,12 @@ def insert(self, **fields):
         if self.conflict_target or self.conflict_action:
             compiler = self._build_insert_compiler([fields])
             rows = compiler.execute_sql(return_id=True)
-            if 'id' in rows[0]:
-                return rows[0]['id']
-            return None
+
+            pk_field_name = self.model._meta.pk.name
+            return rows[0][pk_field_name]
 
         # no special action required, use the standard Django create(..)
-        return super().create(**fields).id
+        return super().create(**fields).pk
 
     def insert_and_get(self, **fields):
         """Creates a new record in the database and then gets

tests/test_insert.py

@@ -0,0 +1,100 @@
+from django.db import models
+
+from psqlextra.query import ConflictAction
+from .util import get_fake_model
+
+
+def test_insert():
+    """Tests whether inserts works when the primary key is explicitly specified."""
+
+    model = get_fake_model({
+        'cookies': models.CharField(max_length=255, null=True),
+    })
+
+    pk = (
+        model.objects.all()
+        .insert(
+            cookies='some-cookies',
+        )
+    )
+
+    assert pk is not None
+
+    obj1 = model.objects.get()
+    assert obj1.pk == pk
+    assert obj1.cookies == 'some-cookies'
+
+
+def test_insert_explicit_pk():
+    """Tests whether inserts works when the primary key is explicitly specified."""
+
+    model = get_fake_model({
+        'name': models.CharField(max_length=255, primary_key=True),
+        'cookies': models.CharField(max_length=255, null=True),
+    })
+
+    pk = (
+        model.objects.all()
+        .insert(
+            name='the-object',
+            cookies='some-cookies',
+        )
+    )
+
+    assert pk == 'the-object'
+
+    obj1 = model.objects.get()
+    assert obj1.pk == 'the-object'
+    assert obj1.name == 'the-object'
+    assert obj1.cookies == 'some-cookies'
+
+
+def test_insert_on_conflict():
+    """Tests whether inserts works when a conflict is anticipated."""
+
+    model = get_fake_model({
+        'name': models.CharField(max_length=255, unique=True),
+        'cookies': models.CharField(max_length=255, null=True),
+    })
+
+    pk = (
+        model.objects.on_conflict([('pk')], ConflictAction.NOTHING)
+        .insert(
+            name='the-object',
+            cookies='some-cookies',
+        )
+    )
+
+    assert pk is not None
+
+    obj1 = model.objects.get()
+    assert obj1.pk == pk
+    assert obj1.name == 'the-object'
+    assert obj1.cookies == 'some-cookies'
+
+
+def test_insert_on_conflict_explicit_pk():
+    """
+    Tests whether inserts works when a conflict is anticipated and the primary
+    key is explicitly specified.
+    """
+
+    model = get_fake_model({
+        'name': models.CharField(max_length=255, primary_key=True),
+        'cookies': models.CharField(max_length=255, null=True),
+    })
+
+    pk = (
+        model.objects.on_conflict([('name')], ConflictAction.NOTHING)
+        .insert(
+            name='the-object',
+            cookies='some-cookies',
+        )
+    )
+
+    assert pk == 'the-object'
+
+    obj1 = model.objects.get()
+    assert obj1.pk == 'the-object'
+    assert obj1.name == 'the-object'
+    assert obj1.cookies == 'some-cookies'

tests/test_upsert.py

@@ -50,6 +50,51 @@ def test_upsert():
     assert obj2.cookies == 'choco'
 
 
+def test_upsert_explicit_pk():
+    """Tests whether upserts works when the primary key is explicitly specified."""
+
+    model = get_fake_model({
+        'name': models.CharField(max_length=255, primary_key=True),
+        'cookies': models.CharField(max_length=255, null=True),
+    })
+
+    obj1 = (
+        model.objects
+        .upsert_and_get(
+            conflict_target=[('name')],
+            fields=dict(
+                name='the-object',
+                cookies='first-cheers',
+            )
+        )
+    )
+
+    obj1.refresh_from_db()
+    assert obj1.name == 'the-object'
+    assert obj1.cookies == 'first-cheers'
+
+    obj2 = (
+        model.objects
+        .upsert_and_get(
+            conflict_target=[('name')],
+            fields=dict(
+                name='the-object',
+                cookies='second-boo',
+            )
+        )
+    )
+
+    obj1.refresh_from_db()
+    obj2.refresh_from_db()
+
+    # assert both objects are the same
+    assert obj1.pk == obj2.pk
+    assert obj1.name == 'the-object'
+    assert obj1.cookies == 'second-boo'
+    assert obj2.name == 'the-object'
+    assert obj2.cookies == 'second-boo'
+
+
 def test_upsert_bulk():
     """Tests whether bulk_upsert works properly."""