Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
AfterFetchQuerySetMixin for Django
class AfterFetchQuerySetMixin:
"""
QuerySet mixin to enable functions to run immediately
after records have been fetched from the DB.
"""
# This is most useful for registering 'prefetch_related' like operations
# or complex aggregations that need to be run after fetching, but while
# still allowing chaining of other QuerySet methods.
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._after_fetch_callbacks = []
def register_after_fetch_callback(self, callback):
"""
Register a callback to be run after the QuerySet is fetched.
The callback should be a callable that accepts a list of model instances.
"""
self._after_fetch_callbacks.append(callback)
return self
# _fetch_all and _clone are Django internals.
def _fetch_all(self):
already_run = self._result_cache is not None
# This super() call fills out the result cache in the QuerySet, and does
# any prefetches.
super()._fetch_all()
if already_run:
# We only run our callbacks once
return
# Now we run our callback.
for c in self._after_fetch_callbacks:
c(self._result_cache)
def _clone(self):
retval = super()._clone()
retval._after_fetch_callbacks = self._after_fetch_callbacks[:]
return retval
# Usage would be like this:
# Example for demo purposes, there are other ways of doing this:
# Suppose we want to decorate each user with an `nth_user_joined`
# attribute which is calculated relative to `date_joined` attribute,
# only for the batch of records retrieved.
class UserQuerySet(AfterFetchQuerySet, models.QuerySet):
def with_nth_user_joined(self):
def add_nth_user_joined(user_list):
for i, user in enumerate(sorted(user_list, key=lambda user: user.date_joined), 1):
user.nth_user_joined = i
return self.register_after_fetch_callback(add_nth_user_joined)
# We can now do the following, with the decoration applied after the query is executed,
# where that query includes the filter.
users = list(User.objects.all().with_nth_user_joined().filter(id__lt=500))
# This technique is especially useful:
# - if you want to do subsequent queries based on the returned values of the first query.
# e.g. things like aggregations or complex prefetches.
# - if you are in some framework where you don't have fully control over when the
# first main query will eventually be executed, but need something to happen immediately
# after that evaluation.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment