Django Recursive Annotation

Question

I'm building a Django app with a recursive comment structure.

Problem: The recursive nature of my comments datastructure means I'm struggling to write a query to annotate each post with the number of replies, and then traverse those posts/replies in my template.

The comment model I've built differentiates between post responses (which are top level comments) and comment responses (which are replies to other comments).

(Post)
3 Total Comments
-----------------
one (post reply)
└── two (comment reply)
    └── three (comment reply)    
(more)

I've represented a comment as follows:

class Comment(TimeStamp):
    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
    content = models.TextField(max_length=2000)
    post = models.ForeignKey("Post", on_delete=models.CASCADE, related_name="comments")
    # Top level comments are those that aren't replies to other comments
    reply = models.ForeignKey(
        "self", on_delete=models.PROTECT, null=True, blank=True, related_name="replies"
    )

This works pretty well, pic related

What works

I'm able to prefetch all the comment replies for a post as follows:

comment_query = Comment.objects.annotate(num_replies=Count("replies"))
post = Post.objects.prefetch_related(Prefetch("comments", comment_query)).get(id="1")

Which correctly displays the number of replies for each comment:

>>> post.comments.values_list('num_replies')                                                                                                                                                 
<QuerySet [(1,), (1,), (0,)]>

What Doesn't Work

This query only annotates the top level post.comments

>>> post.comments.first().replies.all()                                                                                                                                                      
<QuerySet [<Comment: two>]>

>>> post.comments.first().replies.first().num_replies                                                                                                                                       
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-132-8151a7d13021> in <module>
----> 1 post.comments.first().replies.first().num_replies

AttributeError: 'Comment' object has no attribute 'num_replies'                                                                                                                              

In order to render by template properly I need to iterate over comment.replies for each top level response. Any nested comment responses are therefore missing the original num_replies annotation.

In my template/view logic I'm rendering comment trees with roughly the following logic:

{% for comment in post.comments.all %}
{% if not comment.reply %}
  {% include "posts/comment_tree.html" %}
{% endif %}
{% endfor %}

Where post/comments_tree.html contains:

{{ post.content }}
{% for reply in comment.replies.all %}
   {% include "posts/comment_tree.html" with comment=reply %}
{% endfor %}

What I've Tried

I can work around this to an extent by doing the following, which will annotate the first level of replies:

comment_query = Comment.objects.prefetch_related(
    Prefetch("replies", Comment.objects.annotate(num_replies=Count("replies")))
).annotate(num_replies=Count("replies"))

This successfully annotates the second comment, which is a nested response

>>> post.comments.first().replies.first().num_replies                                                                                                                                       
1

But it won't work for any further nested comments (i.e. the third)

>>> post.comments.first().replies.first().replies.first().num_replies                                                                                                                       
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-127-7d5b9798b7b1> in <module>
----> 1 post.comments.first().replies.first().replies.first().num_replies

AttributeError: 'Comment' object has no attribute 'num_replies'

Clearly this approach is completely flawed as I'll be forced to add a nested Prefetch statement for the total number of nested comments I want to support. Ideally I'd like a solution which allows me to annotate an abitrarily nested (self referential) data structure.

TLDR: Is this type of query even possible in Django's ORM, or will I have to reach for SQL?

Answer 1

Take a look at django-cte. You want to define a CTE (common table expression) that contains the annotate. Then use that CTE in the query fetching the comments for a post.

From django-cte's docs:

class Region(Model):
    objects = CTEManager()
    name = TextField(primary_key=True)
    parent = ForeignKey("self", null=True, on_delete=CASCADE)

def make_regions_cte(cte):
    return Region.objects.filter(
        # start with root nodes
        parent__isnull=True
    ).values(
        "name",
        path=F("name"),
        depth=Value(0, output_field=IntegerField()),
    ).union(
        # recursive union: get descendants
        cte.join(Region, parent=cte.col.name).values(
            "name",
            path=Concat(
                cte.col.path, Value("\x01"), F("name"),
                output_field=TextField(),
            ),
            depth=cte.col.depth + Value(1, output_field=IntegerField()),
        ),
        all=True,
    )

cte = With.recursive(make_regions_cte)

regions = (
    cte.join(Region, name=cte.col.name)
    .with_cte(cte)
    .annotate(
        path=cte.col.path,
        depth=cte.col.depth,
    )
    .order_by("path")
)