[Django]-How to prefetch descendants with django-treebeard's MP_Node?

by

3👍

✅

In the end I did end up with adding a foreign key to each task pointing to its root node like such:

root_node = models.ForeignKey('self', null=True,
                          related_name='descendant_tasks',
                          verbose_name=_('root task')
                          )

I updated my save method on my Task model to make sure I always point to the correct root node 

def save(self, force_insert=False, force_update=False, using=None, update_fields=None):
    try:
        self.root_task = self.get_root()
    except ObjectDoesNotExist:
        self.root_task = None

    return super(Task, self).save(force_insert=False, force_update=False, using=None,
                                  update_fields=None
                                  )

and this allows me to simply prefetch all descendants using prefetch_related('descendants').

Whenever I need to have the descendants in a nested fashion I use the following function to nest the flattened list of descendants again

def build_nested(tasks):

    def get_basepath(path, depth):
        return path[0:depth * Task.steplen]

    container, link = [], {}
    for task in sorted(tasks, key=attrgetter('depth')):
        depth = int(len(task.path) / Task.steplen)
        try:
            parent_path = get_basepath(task.path, depth - 1)
            parent_obj = link[parent_path]
            if not hasattr(parent_obj, 'sub_tasks'):
                parent_obj.sub_tasks = []
            parent_obj.sub_tasks.append(task)
        except KeyError:  # Append it as root task if no parent exists
            container.append(task)

        link[task.path] = task

    return container
👤Matthijs Brouns

1👍

If you want to avoid using a Foreign Key you can iterate over the queryset and re-create the tree structure in memory.

In my case I wanted to have a template tag (much like django-mptt’s recursetree templatetag) to show multiple levels of nested pages with only one database query. Basically copying mptt.utils.get_cached_trees I ended up with this:

def get_cached_trees(queryset: QuerySet) -> list:
    """Return top-most pages / roots.

    Each page will have its children stored in `_cached_children` attribute
    and its parent in `_cached_parent`. This avoids having to query the database.
    """
    top_nodes: list = []
    path: list = []
    for obj in queryset:
        obj._cached_children = []
        if obj.depth == queryset[0].depth:
            add_top_node(obj, top_nodes, path)
        else:
            while not is_child_of(obj, parent := path[-1]):
                path.pop()
            add_child(parent, obj)

        if obj.numchild:
            path.append(obj)

    return top_nodes

def add_top_node(obj: MP_Node, top_nodes: list, path: list) -> None:
    top_nodes.append(obj)
    path.clear()

def add_child(parent: MP_Node, obj: MP_Node) -> None:
    obj._cached_parent = parent
    parent._cached_children.append(obj)

def is_child_of(child: MP_Node, parent: MP_Node) -> bool:
    """Return whether `child` is a sub page of `parent` without database query.

    `_get_children_path_interval` is an internal method of MP_Node.
    """
    start, end = parent._get_children_path_interval(parent.path)
    return start < child.path < end

It can be used like this to avoid the dreaded N+1 query problem:

for page in get_cached_trees(queryset):
    for child in page._cached_children:
        ...
👤jnns

Leave a comment