Add support for atomic upserts

It feels like making update_fields also accept a dict[str, Expression] (today it only accepts list[str]) would be the most straight forward way to implement this feature.

Today doing

AbTestHourlyLog.objects.bulk_create([
    AbTestHourlyLog(...),
    update_conflicts=True,
    unique_fields=["ab_test", "version", "date", "hour"],
    update_fields=["participants", "conversions"],
])

Results in something like

INSERT INTO abtesthourlylog (ab_test_id, version, date, hour, participants, conversions)
VALUES (%s, %s, %s, %s, %s, %s)
ON CONFLICT (ab_test_id, version, date, hour)
  DO UPDATE SET participants = EXCLUDED.participants, conversions = EXCLUDED.conversions;

Looking at how bulk_create, SQLInsertLogic, and on_conflict_suffix_sql are implemented is seems relatively straightforward to add support for

AbTestHourlyLog.objects.bulk_create([
    AbTestHourlyLog(...),
    update_conflicts=True,
    unique_fields=["ab_test", "version", "date", "hour"],
    update_fields={
        "participants": F("participants") + participants_incr,
        "conversions": F("conversions") + conversions_incr,
    }
])

To result in

INSERT INTO abtesthourlylog (ab_test_id, version, date, hour, participants, conversions)
VALUES (%s, %s, %s, %s, %s, %s)
ON CONFLICT (ab_test_id, version, date, hour)
  DO UPDATE SET participants = EXCLUDED.participants + %s, conversions = EXCLUDED.conversions + %s;

In order to avoid peppering the logic with isinstance(update_fields, list) or dict we could immediately turn list[str] into dict[str, F] in QuerySet.bulk_create to keep it as a quick alias and more importantly preserve backward compatiblity

if isinstance(update_fields, list):
    update_fields = {
        update_field: F(update_field)
        for update_field in update_fields
    }

All the downstream logic would then only have to expect update_fields: dict[str, Expression] | None.

4 Likes