djangodjango-rest-frameworkdjango-rest-viewsets

Django Rest Framework Filtering an object in get_queryset method


Basically, I have a catalog viewset. In the list view I want to make a few filtering and return accordingly. Relevant Catalog model fields are:

class Catalog(models.Model):
    name = models.CharField(max_length=191, null=True, blank=False)
    ...
    team = models.ForeignKey(Team, on_delete=models.CASCADE, editable=False, related_name='catalogs')
    whitelist_users = models.JSONField(null=True, blank=True, default=list) # If white list is null, it is open to whole team

 

Views.py

class CatalogViewSet(viewsets.ModelViewSet):
    permission_classes = (IsOwnerAdminOrRestricted,)

    def get_queryset(self):
        result = []
        user = self.request.user
        catalogs = Catalog.objects.filter(team__in=self.request.user.team_set.all())
        for catalog in catalogs:
            if catalog.whitelist_users == [] or catalog.whitelist_users == None:
                # catalog is open to whole team
                result.append(catalog)
            else:
                # catalog is private
                if user in catalog.whitelist_users:
                    result.append(catalog)
        return result

So this is my logic;

1 - Get the catalog object if catalog's team is one of the current user' team.

2 - Check if the catalog.whitelist_users contains the current user. (There is also an exception that if it is none means it s open to whole team so I can show it in the list view.)

Now this worked but since I am returning an array, it doesn't find the detail objects correctly. I mean /catalog/ID doesn't work correctly.

I am new to DRF so I am guessing there is something wrong here. How would you implement this filtering better?


Solution

  • As the name of the method suggests, you need to return a queryset. Also, avoid iterating over a queryset if that's not necessary. It's better to do it in a single database hit. For complex queries, you can use the Q object.

    from django.db.models import Q
    
    # ...
    
        def get_queryset(self):
            user = self.request.user
            catalogs = Catalog.objects.filter(
                       Q(whitelist_users__in=[None, []]) | Q(whitelist_users__contains=user),
                       team__in=user.team_set.all())
                                                            
            return catalogs
    

    Now I am not 100% sure the whitelist_users__contains=user will work since it depends on how you construct your JSON, but the idea is there, you will just need to adapt what it contains.

    This will be much more effective than looping in python and will respect what get_queryset is meant for.