Make .upsert() work with one-to-one fields or custom primary key names

Author: Photonios

psqlextra/query.py

@@ -196,11 +196,11 @@ def insert(self, using: Optional[str] = None, **fields):
             compiler = self._build_insert_compiler([fields], using=using)
             rows = compiler.execute_sql(return_id=True)
 
-            pk_field_name = self.model._meta.pk.name
+            _, pk_db_column = self.model._meta.pk.get_attname_column()
             if not rows or len(rows) == 0:
                 return None
 
-            return rows[0][pk_field_name]
+            return rows[0][pk_db_column]
 
         # no special action required, use the standard Django create(..)
         return super().create(**fields).pk

tests/test_upsert.py

@@ -82,6 +82,22 @@ def test_upsert_explicit_pk():
     assert obj2.cookies == "second-boo"
 
 
+def test_upsert_one_to_one_field():
+    model1 = get_fake_model({"title": models.TextField(unique=True)})
+    model2 = get_fake_model(
+        {"model1": models.OneToOneField(model1, on_delete=models.CASCADE)}
+    )
+
+    obj1 = model1.objects.create(title="hello world")
+
+    obj2_id = model2.objects.upsert(
+        conflict_target=["model1"], fields=dict(model1=obj1)
+    )
+
+    obj2 = model2.objects.get(id=obj2_id)
+    assert obj2.model1 == obj1
+
+
 def test_upsert_with_update_condition():
     """Tests that an expression can be used as an upsert update condition."""