Django Dry approach to get_queryset

Question

My app maps users to companies, so they can only see their companies data unless the user is_staff, which can see all.

I am writing the same query over and over again and would like to refactor it to use the dry approach.

The problem is that nothing gets returned. No error either

models.py

class Company(models.Model):
    ...

class Site(models.Model):
    company = models.ForeignKey(Company)
    ...

class UserProfile(AbstractUser):
    company = models.ForeignKey(Company)
    ...

views.py

def get_company(user):
    if user.is_staff:
        company = Company.objects.all().order_by('-name')

    else:
        company = Company.objects.get(id=user.company.id).order_by('-name')

    return company


class DashboardList(ListView):
    model = Company
    template_name = 'sites/dashboard.html'
    paginate_by = 25

    def get_queryset(self):
        return get_company(self.request.user)

Answer 1

Company.objects.get(id=user.company.id).order_by('-name')

Is not a valid operation. get returns a company object, not a queryset

Try filter instead of get

Company.objects.filter(id=user.company.id).order_by('-name')

It will be fine. So you can refactor like this

def get_company(user):
    query = {}
    if not user.is_staff:
        query['id'] = user.company.id
    return Company.objects.filter(**query).order_by('-name')

Answer 2

Since you want to map userS to companIES you will need to use a many to many relationship, instead of a one to many. Now you are just stating that a user can have only one company associated, but a company can have many users associated. So doing a filter for a user's companies makes no sense. To fix it, your models.py should be:

class Company(models.Model):
    ...

class UserProfile(AbstractUser):
    companies = models.ManyToManyField(Company)
    ...

and your view.py:

def get_companies(user):
    if user.is_staff:
        companies = Company.objects.all()
    else:
        companies = user.companies.all()

    return companies.order_by('-name')


class DashboardList(ListView):
    model = Company
    template_name = 'sites/dashboard.html'
    paginate_by = 25

    def get_queryset(self):
        return get_company(self.request.user)