Optimal way to fetch related model and through model together in ManyToManyField-CodePudding

Given schema:

class Investor(models.Model):
    name = models.CharField(max_length=30)
    advisors = models.ManyToManyField("Advisor", related_name="investors", through="Connection")

class Advisor(models.Model):
    name = models.CharField(max_length=30)

class Connection(models.Model):
    investor = models.ForeignKey("Investor", related_name="connections", on_delete=models.CASCADE)
    advisor  = models.ForeignKey("Advisor",  related_name="connections", on_delete=models.CASCADE)
    blocked  = models.BooleanField(default=False)

And given an Advisor "a", what is the optimal to fetch a list of all Investors and their Connection to "a" when it exists?

The best I've figured out so far is:

from django.db.models import Prefetch

for investor in Investor.objects.prefetch_related(
    Prefetch(
        'connections',
        queryset=Connection.objects.filter(advisor=a),
        to_attr='connection',
    )
):
    name = investor.name
    blocked = None
    if investor.connection:
        blocked = investor.connection[0].blocked
    print(f"{name} (blocked={blocked})")

Update

I found a way to make using the above strategy more pleasant:

class WithConnection:
    def __init__(self, queryset):
        self.queryset = queryset

    def __iter__(self):
        for model in self.queryset:
            connection = model.my_connections[0] if model.my_connections else None
            yield model, connection

class InvestorQuerySet(models.QuerySet):
    def pov(self, advisor):
        return WithConnection(
            self.prefetch_related(models.Prefetch(
                'connections',
                queryset=Connection.objects.filter(advisor=advisor),
                to_attr='my_connections',
            ))
        )

Used like this:

advisor = Advisor.objects.first()
for investor, connection in Investor.objects.pov(advisor):
    # do stuff with investor
    if connection:
        # do stuff with connection

Would like to improve on this. Django's ORM is so powerful, seems like there ought to be a simpler way to achieve this.

CodePudding user response：

I beleive you simply want

a.investor_set.all()

CodePudding user response：

After much research and many iterations, I’ve arrived at this solution:

class WithConnectionQuerySet(models.QuerySet):
    def __iter__(self):
        return map(
            lambda m: (m, m.my_connections[0]) if m.my_connections else (m, None),
            super().__iter__(),
        )

class Investor(models.Model):
    ...
    @classmethod
    def viewed_by(cls, advisor):
        manager = cls._default_manager.__class__.from_queryset(WithConnectionQuerySet)()
        manager.model = cls
        return manager.prefetch_related(models.Prefetch(
            'connections',
            queryset=Connection.objects.filter(advisor=advisor),
            to_attr='my_connections',
        ))

Used like this:

advisor = Advisor.objects.first()
for investor, connection in Investor.viewed_by(advisor).filter(age__gt=40):
    # do stuff with investor
    if connection:
        # do stuff with connection

I'm not still not 100% thrilled with it, but I like it enough now to start using it. If anyone has a better answer, please chime in.