comment_prefetch.py

from db.models import Comment, User
from django.db.models import Prefetch
from django_print_sql import print_sql

import rich
import uuid

class User(models.Model):
    name = models.CharField(max_length=100)

    def __str__(self):
        return self.name


class Comment(models.Model):
    user = models.ForeignKey(User, related_name="comments", on_delete=models.CASCADE)
    parent = models.ForeignKey("self", related_name="replies", null=True, blank=True, on_delete=models.CASCADE)
    created_at = models.DateTimeField(auto_now_add=True)
    replies: "RelatedManager[Comment]"
    content = models.TextField()

c = Comment.objects.create(user=User.objects.create(name=uuid.uuid4().hex), content="This is a top-level comment.")
c1 = Comment.objects.create(user=User.objects.create(name=uuid.uuid4().hex), parent=c, content="This is a reply to the top-level comment.")
c2 = Comment.objects.create(user=User.objects.create(name=uuid.uuid4().hex), parent=c, content="This is another reply to the top-level comment.")
c3 = Comment.objects.create(user=User.objects.create(name=uuid.uuid4().hex), parent=c, content="This is a reply to the first reply.")
cc = Comment.objects.create(user=User.objects.create(name=uuid.uuid4().hex), content="This is a top-level comment with no replies.")
cc1 = Comment.objects.create(user=User.objects.create(name=uuid.uuid4().hex), parent=cc, content="This is a reply to the top-level comment with no replies.")
cc2 = Comment.objects.create(
    user=User.objects.create(name=uuid.uuid4().hex), parent=cc, content="This is another reply to the top-level comment with no replies."
)
cc3 = Comment.objects.create(
    user=User.objects.create(name=uuid.uuid4().hex), parent=cc, content="This is a reply to the first reply to the top-level comment with no replies."
)

with print_sql():
    data = []
    top_level_comments = (
        Comment.objects.filter(parent__isnull=True)
        .prefetch_related(
            Prefetch("replies", queryset=Comment.objects.order_by("created_at")[:2], to_attr="top_replies"),
            Prefetch("top_replies__user"),
            Prefetch("user"),
        )
        .annotate(
            num_replies=Count("replies"),
        )
    )

    def get_top_replies(comment) -> list[Comment]:
        return getattr(comment, "top_replies", [])

    def get_num_replies(comment) -> int:
        return getattr(comment, "num_replies", 0)

    for comment in top_level_comments:
        data.append(
            {
                "comment.id": comment.id,
                "comment.user.name": comment.user.name,
                "comment.content": comment.content,
                "comment.num_replies": get_num_replies(comment),
                "comment.replies": [
                    {
                        "reply.id": reply.id,
                        "reply.user.name": reply.user.name,
                        "reply.content": reply.content,
                    }
                    for reply in get_top_replies(comment)
                ],
            }
        )
    rich.print(data)

# [
#     {
#         'comment.id': 1,
#         'comment.user.name': 'af73bfb482024c079c288e12a09545d6',
#         'comment.content': 'This is a top-level comment.',
#         'comment.num_replies': 3,
#         'comment.replies': [
#             {'reply.id': 2, 'reply.user.name': '741eff99fb734ce388994441d50b8ff5', 'reply.content': 'This is a reply to the top-level comment.'},
#             {'reply.id': 3, 'reply.user.name': 'd836ab05dbd6449cb538fb5c336d4210', 'reply.content': 'This is another reply to the top-level comment.'}
#         ]
#     },
#     {
#         'comment.id': 5,
#         'comment.user.name': 'c96caa6990a24341ae029105572453c1',
#         'comment.content': 'This is a top-level comment with no replies.',
#         'comment.num_replies': 3,
#         'comment.replies': [
#             {'reply.id': 6, 'reply.user.name': '9975b0b66cd44b4aa1948b98ed596630', 'reply.content': 'This is a reply to the top-level comment with no replies.'},    
#             {
#                 'reply.id': 7,
#                 'reply.user.name': 'b639c899a5db44689987f24a9f0a7248',
#                 'reply.content': 'This is another reply to the top-level comment with no replies.'
#             }
#         ]
#     }
# ]
# [4 queries executed, total time elapsed 4.02ms]