Django's JSONField incrementation with F expressions

Vasil Slavov
Sep 16, 2022
Categories:Django

Updating a QuerySet using an F expression is a common and very useful Django pattern for updating a set of objects with a reduced number of queries, while having concurrency in mind.

We're often using it to increment stored counters in our models like that -

SomeDataModel.objects.filter(some_filter).update(stored_field=F("stored_field") + 1)

However, the above pattern does not work if we're trying to apply it to keys in a JSONField.

The problem

Let's say we have several stored counters in a model and we're keeping them in a JSONField in order to add / remove counters dynamically without creating additional fields & migrations.

class SomeDataModel(models.Model):
    stored_field = models.JSONField(blank=True, default=dict)


SomeDataModel.objects.create(
    stored_field = {
        "first_key": 0,
        "second_key": 0
    }
)

If we want to query our model, we can simply chain the field and the key name to form a path:

In: SomeDataModel.objects.filter(stored_field__first_key=10)

Out: <QuerySet [<SomeDataModel: 1>]>

Now, if we try updating our queryset using the same key we'll encounter an error:

SomeDataModel.objects.filter(
    stored_field__first_key=10
).update(
    stored_field__first_key=F('stored_field__first_key') + 1
)

Results in:

FieldDoesNotExist: SomeDataModel has no field named 'stored_field__first_key'

The problem here is that JSONField key and path transforms are not supported in QuerySet.update().

There is an active Django ticket which will allow us to do much more with JSONFields once it's completed.

The solution

Since we really like the interface of the F expression update, we started looking for a solution.

What we ended up doing was:

  1. Find out how to update a json field using SQL (props to this Stack Overflow answer.)
  2. Use a RawSQL expression to update our QuerySet.

And the initial version of it was:

from django.db.models.expressions import RawSQL


def jsonfield_increment(field_name: str, key: str, increment_by: int) -> RawSQL:
    sql = f"""
        jsonb_set(
            {field_name},
            '{{{key}}}',
            (COALESCE({field_name}->>'{key}','0')::int + {increment_by})::text::jsonb
        )
    """

    return RawSQL(sql, [])

# Usage
SomeDataModel.objects.filter(
    stored_field__first_key=10
).update(
    stored_field=jsonfield_increment("stored_field", "first_key", 1)
)

Let's explain what's happening here:

  1. jsonb_set is a database function which returns a new jsonb value but with a different value for a given key.
  2. We're specifying that the jsonb value we want to change is contained in the field_name column and the key we want to change.
  3. We're specifying that the new value in this key should be:

4. The return value from this will be a new jsonb value which we're writing inside the stored_field column.

We can leave the code above as it is and it will do the job.

If we want to remove the raw SQL from our codebase, we can iterate a bit further.

The good thing is all of the database functions we're using are already available as custom Django ORM expressions or functions.

Here's how we can update our Queryset when we're using the custom ORM expressions:

from django.db.models import F, Value, IntegerField, CharField, JSONField
from django.db.models.functions import Cast
from django.db.models.expressions import Func


SomeDataModel.objects.filter(stored_field__first_key=10).update(
    field=Func(
        F('stored_field'),
        Value('{first_key}'),
        Cast(
            Cast(
                Cast(
                    F('stored_field__first_key'),
                    IntegerField()
                ) + 1,
                CharField()
            ),
            JSONField()
        ),
        function="jsonb_set"
    )
)

Of course, we don't want to repeat all that every time we need to increment a key inside a JSONField, so we can wrap it inside a Func subclass:

from django.db.models import F, Value, IntegerField, CharField, JSONField
from django.db.models.functions import Cast
from django.db.models.expressions import Func


class JSONIncrement(Func):
    function = "jsonb_set"

    def __init__(self, full_path, value=1, **extra):
        field_name, *key_path_parts = full_path.split("__")

        if not field_name:
            raise ValueError("`full_path` can not be blank.")

        if len(key_path_parts) < 1:
            raise ValueError("`full_path` must contain at least one key.")

        key_path = ",".join(key_path_parts)

        new_value_expr = Cast(
            Cast(F(full_path), IntegerField()) + value,
            CharField()
        )

        expressions = [
            F(field_name),
            Value(f"{{{key_path}}}"),
            Cast(new_value_expr, JSONField())
        ]

        super().__init__(*expressions, output_field=JSONField(), **extra)

And the usage would look like this:

# Usage
SomeDataModel.objects.filter(
    stored_field__first_key=10
).update(
    JSONIncrement(stored_field__first_key, value=1)
)

And that's it. We achieved out desired F expression-like interface, for incrementing counters in a JSONField, in Django.

Hopefully, you'll find this useful 🙌