Move problem filtering into class methods (#1308)

judge/feed.py

@@ -15,7 +15,7 @@ class ProblemFeed(Feed):
     description = 'The latest problems added on the %s website' % settings.SITE_LONG_NAME
 
     def items(self):
-        return Problem.objects.filter(is_public=True, is_organization_private=False).order_by('-date', '-id')[:25]
+        return Problem.get_public_problems().order_by('-date', '-id')[:25]
 
     def item_title(self, problem):
         return problem.name

judge/models/problem.py

@@ -5,7 +5,7 @@ from django.contrib.contenttypes.fields import GenericRelation
 from django.core.cache import cache
 from django.core.validators import MaxValueValidator, MinValueValidator, RegexValidator
 from django.db import models
-from django.db.models import CASCADE, F, QuerySet, SET_NULL
+from django.db.models import CASCADE, F, Q, QuerySet, SET_NULL
 from django.db.models.expressions import RawSQL
 from django.db.models.functions import Coalesce
 from django.urls import reverse
@@ -224,6 +224,43 @@ class Problem(models.Model):
     def is_subs_manageable_by(self, user):
         return user.is_staff and user.has_perm('judge.rejudge_submission') and self.is_editable_by(user)
 
+    @classmethod
+    def get_visible_problems(cls, user):
+        # Do unauthenticated check here so we can skip authentication checks later on.
+        if not user.is_authenticated:
+            return cls.get_public_problems()
+
+        # Conditions for visible problem:
+        #   - `judge.edit_all_problem` or `judge.see_private_problem`
+        #   - otherwise
+        #       - not is_public problems
+        #           - author or curator or tester
+        #       - is_public problems
+        #           - not is_organization_private or in organization or `judge.see_organization_problem`
+        #           - author or curator or tester
+        queryset = cls.objects.defer('description')
+
+        if not (user.has_perm('judge.see_private_problem') or user.has_perm('judge.edit_all_problem')):
+            q = Q(is_public=True)
+            if not user.has_perm('judge.see_organization_problem'):
+                # Either not organization private or in the organization.
+                q &= (
+                    Q(is_organization_private=False) |
+                    Q(is_organization_private=True, organizations__in=user.profile.organizations.all())
+                )
+
+            # Authors, curators, and testers should always have access, so OR at the very end.
+            q |= Q(authors=user.profile)
+            q |= Q(curators=user.profile)
+            q |= Q(testers=user.profile)
+            queryset = queryset.filter(q)
+
+        return queryset.distinct()
+
+    @classmethod
+    def get_public_problems(cls):
+        return cls.objects.filter(is_public=True, is_organization_private=False).defer('description').distinct()
+
     def __str__(self):
         return self.name
 

judge/models/profile.py

@@ -138,12 +138,15 @@ class Profile(models.Model):
 
     def calculate_points(self, table=_pp_table):
         from judge.models import Problem
-        data = (Problem.objects.filter(submission__user=self, submission__points__isnull=False, is_public=True,
-                                       is_organization_private=False)
-                       .annotate(max_points=Max('submission__points')).order_by('-max_points')
-                       .values_list('max_points', flat=True).filter(max_points__gt=0))
-        extradata = Problem.objects.filter(submission__user=self, submission__result='AC', is_public=True) \
-                           .values('id').distinct().count()
+        public_problems = Problem.get_public_problems()
+        data = (
+            public_problems.filter(submission__user=self, submission__points__isnull=False)
+                           .annotate(max_points=Max('submission__points')).order_by('-max_points')
+                           .values_list('max_points', flat=True).filter(max_points__gt=0)
+        )
+        extradata = (
+            public_problems.filter(submission__user=self, submission__result='AC').values('id').count()
+        )
         bonus_function = settings.DMOJ_PP_BONUS_FUNCTION
         points = sum(data)
         problems = len(data)

judge/sitemap.py

@@ -11,7 +11,7 @@ class ProblemSitemap(Sitemap):
     priority = 0.8
 
     def items(self):
-        return Problem.objects.filter(is_public=True, is_organization_private=False).values_list('code')
+        return Problem.get_public_problems().values_list('code')
 
     def location(self, obj):
         return reverse('problem_detail', args=obj)

judge/utils/problems.py

@@ -112,8 +112,8 @@ def hot_problems(duration, limit):
     cache_key = 'hot_problems:%d:%d' % (duration.total_seconds(), limit)
     qs = cache.get(cache_key)
     if qs is None:
-        qs = Problem.objects.filter(is_public=True, is_organization_private=False,
-                                    submission__date__gt=timezone.now() - duration, points__gt=3, points__lt=25)
+        qs = Problem.get_public_problems() \
+                    .filter(submission__date__gt=timezone.now() - duration, points__gt=3, points__lt=25)
         qs0 = qs.annotate(k=Count('submission__user', distinct=True)).order_by('-k').values_list('k', flat=True)
 
         if not qs0:

judge/views/api/api_v1.py

@@ -90,7 +90,7 @@ def api_v1_contest_detail(request, contest):
 
 
 def api_v1_problem_list(request):
-    queryset = Problem.objects.filter(is_public=True, is_organization_private=False)
+    queryset = Problem.get_visible_problems(request.user)
     if settings.ENABLE_FTS and 'search' in request.GET:
         query = ' '.join(request.GET.getlist('search')).strip()
         if query:

judge/views/blog.py

@@ -39,7 +39,7 @@ class PostList(ListView):
         context['first_page_href'] = reverse('home')
         context['page_prefix'] = reverse('blog_post_list')
         context['comments'] = Comment.most_recent(self.request.user, 10)
-        context['new_problems'] = Problem.objects.filter(is_public=True, is_organization_private=False) \
+        context['new_problems'] = Problem.get_public_problems() \
                                          .order_by('-date', '-id')[:settings.DMOJ_BLOG_NEW_PROBLEM_COUNT]
         context['page_titles'] = CacheDict(lambda page: Comment.get_page_title(page))
 
@@ -52,7 +52,7 @@ class PostList(ListView):
                 context['clarifications'] = clarifications.order_by('-date')
 
         context['user_count'] = lazy(Profile.objects.count, int, int)
-        context['problem_count'] = lazy(Problem.objects.filter(is_public=True).count, int, int)
+        context['problem_count'] = lazy(Problem.get_public_problems().count, int, int)
         context['submission_count'] = lazy(Submission.objects.count, int, int)
         context['language_count'] = lazy(Language.objects.count, int, int)
 

judge/views/select2.py

@@ -54,13 +54,8 @@ class OrganizationSelect2View(Select2View):
 
 class ProblemSelect2View(Select2View):
     def get_queryset(self):
-        queryset = Problem.objects.filter(Q(code__icontains=self.term) | Q(name__icontains=self.term))
-        if not self.request.user.has_perm('judge.see_private_problem'):
-            filter = Q(is_public=True)
-            if self.request.user.is_authenticated:
-                filter |= Q(authors=self.request.profile) | Q(curators=self.request.profile)
-            queryset = queryset.filter(filter).distinct()
-        return queryset.distinct()
+        return Problem.get_visible_problems(self.request.user) \
+                      .filter(Q(code__icontains=self.term) | Q(name__icontains=self.term))
 
 
 class ContestSelect2View(Select2View):

judge/views/submission.py

@@ -254,13 +254,7 @@ class SubmissionsListBase(DiggPaginatorMixin, TitleMixin, ListView):
     def get_queryset(self):
         queryset = self._get_queryset()
         if not self.in_contest:
-            if not self.request.user.has_perm('judge.see_private_problem'):
-                queryset = queryset.filter(problem__is_public=True)
-            if not self.request.user.has_perm('judge.see_organization_problem'):
-                filter = Q(problem__is_organization_private=False)
-                if self.request.user.is_authenticated:
-                    filter |= Q(problem__organizations__in=self.request.profile.organizations.all())
-                queryset = queryset.filter(filter)
+            queryset = queryset.filter(problem__in=Problem.get_visible_problems(self.request.user))
         return queryset
 
     def get_my_submissions_page(self):