Skip to content

Commit 74af280

Browse files
committed
Make .upsert() work with one-to-one fields or custom primary key names
1 parent 8b24323 commit 74af280

File tree

2 files changed

+18
-2
lines changed

2 files changed

+18
-2
lines changed

psqlextra/query.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -196,11 +196,11 @@ def insert(self, using: Optional[str] = None, **fields):
196196
compiler = self._build_insert_compiler([fields], using=using)
197197
rows = compiler.execute_sql(return_id=True)
198198

199-
pk_field_name = self.model._meta.pk.name
199+
_, pk_db_column = self.model._meta.pk.get_attname_column()
200200
if not rows or len(rows) == 0:
201201
return None
202202

203-
return rows[0][pk_field_name]
203+
return rows[0][pk_db_column]
204204

205205
# no special action required, use the standard Django create(..)
206206
return super().create(**fields).pk

tests/test_upsert.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,22 @@ def test_upsert_explicit_pk():
8282
assert obj2.cookies == "second-boo"
8383

8484

85+
def test_upsert_one_to_one_field():
86+
model1 = get_fake_model({"title": models.TextField(unique=True)})
87+
model2 = get_fake_model(
88+
{"model1": models.OneToOneField(model1, on_delete=models.CASCADE)}
89+
)
90+
91+
obj1 = model1.objects.create(title="hello world")
92+
93+
obj2_id = model2.objects.upsert(
94+
conflict_target=["model1"], fields=dict(model1=obj1)
95+
)
96+
97+
obj2 = model2.objects.get(id=obj2_id)
98+
assert obj2.model1 == obj1
99+
100+
85101
def test_upsert_with_update_condition():
86102
"""Tests that an expression can be used as an upsert update condition."""
87103

0 commit comments

Comments
 (0)