how to filter nested related django objects

Question

I have an app with lots of investors that invest in the same rounds, which belong to companies, as seen below. However when a user(investor) is logged in, i only want him to be able to see HIS investments.

    {
        "id": 1,
        "name": "Technology Company",
        "rounds": [
            {
                "id": 1,
                "kind": "priced round",
                "company": 1,
                "investments": [
                    {
                        "id": 1,
                        "investor": 1,
                        "round": 1,
                        "size": 118000,
                    },
                    {
                        "id": 2,
                        "investor": 2,
                        "round": 1,
                        "size": 183000,
                    },
                ]
            }
        ]
    },

Currently, my viewsets extend get_queryset as so:

class CompanyViewSet(viewsets.ModelViewSet):
    def get_queryset(self):
        user = self.request.user
        investor = Investor.objects.get(user=user)
        companies = Company.objects.filter(rounds__investments__investor=investor)
        return companies

It retrieves the investments that belong to the investor, but when it takes those investments again to retrieve the rounds, it grabs the round with ALL investors.

How can i write this such that it only displays the investments that below to the Investor?

Here are my models:

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

class Round(PolymorphicModel):
    company = models.ForeignKey(Company, related_name='rounds', blank=True, null=True)

class Investment(PolymorphicModel):
    investor = models.ForeignKey(Investor, related_name='investor')
    size = models.BigIntegerField(default=0)

Solution

Your description of what happens is pretty unclear. What does "when it takes those investments again" mean? Anyway, I'm guessing what you need to do is to use .prefetch_related and a Prefetch object.

from django.db.models import Prefetch

class CompanyViewSet(viewsets.ModelViewSet):
    def get_queryset(self):
        user = self.request.user
        investor = Investor.objects.get(user=user)
        companies = Company.objects.filter(
            rounds__investments__investor_id=investor.id
        ).prefetch_related(Prefetch(
            'rounds__investments',
            queryset=Investment.objects.filter(
                investor_id=investor.pk,
            ),
        ))
        return companies

I haven't tested this snippet but it should give you a pointer in the right direction. I also optimized the investor lookup to check the id only, this will save you an unnecessary indirection.