Merge pull request #46 from PeterJCLaw/fix-foreign-key-on-conflict

Cope with foreign key values being specified as objects

Author: Photonios

psqlextra/compiler.py

@@ -271,7 +271,12 @@ def _format_field_value(self, field_name) -> str:
         return SQLInsertCompiler.prepare_value(
             self,
             field,
-            getattr(self.query.objs[0], field_name)
+            # Note: this deliberately doesn't use `pre_save_val` as we don't
+            # want things like auto_now on DateTimeField (etc.) to change the
+            # value. We rely on pre_save having already been done by the
+            # underlying compiler so that things like FileField have already had
+            # the opportunity to save out their data.
+            getattr(self.query.objs[0], field.attname)
         )
 
     def _normalize_field_name(self, field_name) -> str:

tests/test_on_conflict_nothing.py

@@ -1,3 +1,4 @@
+import pytest
 from django.db import models
 
 from psqlextra.fields import HStoreField
@@ -39,3 +40,112 @@ def test_on_conflict_nothing():
     assert obj1.cookies == 'cheers'
     assert obj2.title['key1'] == 'beer'
     assert obj2.cookies == 'cheers'
+
+
+def test_on_conflict_nothing_foreign_key_by_object():
+    """
+    Tests whether simple insert NOTHING works correctly when the potentially
+    conflicting field is a foreign key specified as an object.
+    """
+
+    other_model = get_fake_model({})
+
+    model = get_fake_model({
+        'other': models.OneToOneField(
+            other_model,
+            on_delete=models.CASCADE,
+        ),
+        'data': models.CharField(max_length=255),
+    })
+
+    other_obj = other_model.objects.create()
+
+    obj1 = (
+        model.objects
+        .on_conflict(['other'], ConflictAction.NOTHING)
+        .insert_and_get(other=other_obj, data="some data")
+    )
+
+    assert obj1.other == other_obj
+    assert obj1.data == "some data"
+
+    obj1.refresh_from_db()
+    assert obj1.other == other_obj
+    assert obj1.data == "some data"
+
+    with pytest.raises(ValueError):
+        (
+            model.objects
+            .on_conflict(['other'], ConflictAction.NOTHING)
+            .insert_and_get(other=obj1)
+        )
+
+    obj2 = (
+        model.objects
+        .on_conflict(['other'], ConflictAction.NOTHING)
+        .insert_and_get(other=other_obj, data="different data")
+    )
+
+    assert obj2.other == other_obj
+    assert obj2.data == "some data"
+
+    obj1.refresh_from_db()
+    obj2.refresh_from_db()
+
+    # assert that the 'other' field didn't change
+    assert obj1.id == obj2.id
+    assert obj1.other == other_obj
+    assert obj2.other == other_obj
+    assert obj1.data == "some data"
+    assert obj2.data == "some data"
+
+
+def test_on_conflict_nothing_foreign_key_by_id():
+    """
+    Tests whether simple insert NOTHING works correctly when the potentially
+    conflicting field is a foreign key specified as an id.
+    """
+
+    other_model = get_fake_model({})
+
+    model = get_fake_model({
+        'other': models.OneToOneField(
+            other_model,
+            on_delete=models.CASCADE,
+        ),
+        'data': models.CharField(max_length=255),
+    })
+
+    other_obj = other_model.objects.create()
+
+    obj1 = (
+        model.objects
+        .on_conflict(['other_id'], ConflictAction.NOTHING)
+        .insert_and_get(other_id=other_obj.pk, data="some data")
+    )
+
+    assert obj1.other == other_obj
+    assert obj1.data == "some data"
+
+    obj1.refresh_from_db()
+    assert obj1.other == other_obj
+    assert obj1.data == "some data"
+
+    obj2 = (
+        model.objects
+        .on_conflict(['other_id'], ConflictAction.NOTHING)
+        .insert_and_get(other_id=other_obj.pk, data="different data")
+    )
+
+    assert obj2.other == other_obj
+    assert obj2.data == "some data"
+
+    obj1.refresh_from_db()
+    obj2.refresh_from_db()
+
+    # assert that the 'other' field didn't change
+    assert obj1.id == obj2.id
+    assert obj1.other == other_obj
+    assert obj2.other == other_obj
+    assert obj1.data == "some data"
+    assert obj2.data == "some data"

tests/test_on_conflict_update.py

@@ -1,3 +1,4 @@
+import pytest
 from django.db import models
 
 from psqlextra.fields import HStoreField
@@ -39,3 +40,112 @@ def test_on_conflict_update():
     assert obj1.cookies == 'choco'
     assert obj2.title['key1'] == 'beer'
     assert obj2.cookies == 'choco'
+
+
+def test_on_conflict_update_foreign_key_by_object():
+    """
+    Tests whether simple upsert works correctly when the conflicting field is a
+    foreign key specified as an object.
+    """
+
+    other_model = get_fake_model({})
+
+    model = get_fake_model({
+        'other': models.OneToOneField(
+            other_model,
+            on_delete=models.CASCADE,
+        ),
+        'data': models.CharField(max_length=255),
+    })
+
+    other_obj = other_model.objects.create()
+
+    obj1 = (
+        model.objects
+        .on_conflict(['other'], ConflictAction.UPDATE)
+        .insert_and_get(other=other_obj, data="some data")
+    )
+
+    assert obj1.other == other_obj
+    assert obj1.data == "some data"
+
+    obj1.refresh_from_db()
+    assert obj1.other == other_obj
+    assert obj1.data == "some data"
+
+    with pytest.raises(ValueError):
+        (
+            model.objects
+            .on_conflict(['other'], ConflictAction.UPDATE)
+            .insert_and_get(other=obj1)
+        )
+
+    obj2 = (
+        model.objects
+        .on_conflict(['other'], ConflictAction.UPDATE)
+        .insert_and_get(other=other_obj, data="different data")
+    )
+
+    assert obj2.other == other_obj
+    assert obj2.data == "different data"
+
+    obj1.refresh_from_db()
+    obj2.refresh_from_db()
+
+    # assert that the 'other' field didn't change
+    assert obj1.id == obj2.id
+    assert obj1.other == other_obj
+    assert obj2.other == other_obj
+    assert obj1.data == "different data"
+    assert obj2.data == "different data"
+
+
+def test_on_conflict_update_foreign_key_by_id():
+    """
+    Tests whether simple upsert works correctly when the conflicting field is a
+    foreign key specified as an id.
+    """
+
+    other_model = get_fake_model({})
+
+    model = get_fake_model({
+        'other': models.OneToOneField(
+            other_model,
+            on_delete=models.CASCADE,
+        ),
+        'data': models.CharField(max_length=255),
+    })
+
+    other_obj = other_model.objects.create()
+
+    obj1 = (
+        model.objects
+        .on_conflict(['other_id'], ConflictAction.UPDATE)
+        .insert_and_get(other_id=other_obj.pk, data="some data")
+    )
+
+    assert obj1.other == other_obj
+    assert obj1.data == "some data"
+
+    obj1.refresh_from_db()
+    assert obj1.other == other_obj
+    assert obj1.data == "some data"
+
+    obj2 = (
+        model.objects
+        .on_conflict(['other_id'], ConflictAction.UPDATE)
+        .insert_and_get(other_id=other_obj.pk, data="different data")
+    )
+
+    assert obj2.other == other_obj
+    assert obj2.data == "different data"
+
+    obj1.refresh_from_db()
+    obj2.refresh_from_db()
+
+    # assert that the 'other' field didn't change
+    assert obj1.id == obj2.id
+    assert obj1.other == other_obj
+    assert obj2.other == other_obj
+    assert obj1.data == "different data"
+    assert obj2.data == "different data"