refactor: migrate contest authors/curators/testers to ContestRole

Author: whlongg

judge/admin/contest.py

@@ -15,6 +15,7 @@
 from reversion.admin import VersionAdmin
 
 from judge.models import Contest, ContestAnnouncement, ContestProblem, ContestSubmission, Profile, Rating, Submission
+from judge.models.role import ContestRole, ROLE_AUTHOR, ROLE_CURATOR
 from judge.ratings import rate_contest
 from judge.utils.views import NoBatchDeleteMixin
 from judge.widgets import AdminAceWidget, AdminHeavySelect2MultipleWidget, AdminHeavySelect2Widget, \
@@ -124,9 +125,6 @@ def clean(self):
 
     class Meta:
         widgets = {
-            'authors': AdminHeavySelect2MultipleWidget(data_view='profile_select2'),
-            'curators': AdminHeavySelect2MultipleWidget(data_view='profile_select2'),
-            'testers': AdminHeavySelect2MultipleWidget(data_view='profile_select2'),
             'private_contestants': AdminHeavySelect2MultipleWidget(data_view='profile_select2'),
             'organization': AdminHeavySelect2Widget(data_view='organization_select2'),
             'tags': AdminSelect2MultipleWidget,
@@ -137,9 +135,23 @@ class Meta:
         }
 
 
+class ContestRoleInlineForm(ModelForm):
+    class Meta:
+        model = ContestRole
+        fields = ('user', 'role')
+        widgets = {'user': AdminHeavySelect2Widget(data_view='profile_select2')}
+
+
+class ContestRoleInline(admin.TabularInline):
+    model = ContestRole
+    fields = ('user', 'role')
+    extra = 0
+    form = ContestRoleInlineForm
+
+
 class ContestAdmin(NoBatchDeleteMixin, SortableAdminBase, VersionAdmin):
     fieldsets = (
-        (None, {'fields': ('key', 'name', 'authors', 'curators', 'testers')}),
+        (None, {'fields': ('key', 'name')}),
         (_('Settings'), {'fields': ('is_visible', 'use_clarifications', 'push_announcements', 'disallow_virtual',
                                     'hide_problem_tags', 'hide_problem_authors', 'show_short_display',
                                     'run_pretests_only', 'locked_after', 'scoreboard_visibility',
@@ -159,7 +171,7 @@ class ContestAdmin(NoBatchDeleteMixin, SortableAdminBase, VersionAdmin):
     list_display = ('key', 'name', 'is_visible', 'is_rated', 'locked_after', 'start_time', 'end_time', 'time_limit',
                     'user_count')
     search_fields = ('key', 'name')
-    inlines = [ContestProblemInline, ContestAnnouncementInline]
+    inlines = [ContestRoleInline, ContestProblemInline, ContestAnnouncementInline]
     actions_on_top = True
     actions_on_bottom = True
     form = ContestForm
@@ -186,7 +198,11 @@ def get_queryset(self, request):
         if request.user.has_perm('judge.edit_all_contest'):
             return queryset
         else:
-            return queryset.filter(Q(authors=request.profile) | Q(curators=request.profile)).distinct()
+            return queryset.filter(
+                contest_roles__user=request.profile,
+            ).filter(
+                Q(contest_roles__role=ROLE_AUTHOR) | Q(contest_roles__role=ROLE_CURATOR),
+            ).distinct()
 
     def get_readonly_fields(self, request, obj=None):
         readonly = []
@@ -358,18 +374,9 @@ def rate_view(self, request, id):
     def get_form(self, request, obj=None, **kwargs):
         form = super(ContestAdmin, self).get_form(request, obj, **kwargs)
         if 'problem_label_script' in form.base_fields:
-            # form.base_fields['problem_label_script'] does not exist when the user has only view permission
-            # on the model.
             form.base_fields['problem_label_script'].widget = AdminAceWidget(
                 mode='lua', theme=request.profile.resolved_ace_theme,
             )
-
-        perms = ('edit_own_contest', 'edit_all_contest')
-        form.base_fields['curators'].queryset = Profile.objects.filter(
-            Q(user__is_superuser=True) |
-            Q(user__groups__permissions__codename__in=perms) |
-            Q(user__user_permissions__codename__in=perms),
-        ).distinct()
         return form
 
 

judge/jinja2/submission.py

@@ -1,12 +1,12 @@
-from operator import attrgetter
-
 from judge.models import SubmissionSourceAccess
+from judge.models.role import ContestRole, ROLE_AUTHOR, ROLE_CURATOR
 from . import registry
 
 
-# TODO: maybe refactor this?
 def get_editor_ids(contest):
-    return set(map(attrgetter('id'), contest.authors.all())) | set(map(attrgetter('id'), contest.curators.all()))
+    return set(ContestRole.objects.filter(
+        contest=contest, role__in=[ROLE_AUTHOR, ROLE_CURATOR],
+    ).values_list('user_id', flat=True))
 
 
 @registry.function

judge/migrations/0220_migrate_contest_roles.py

@@ -0,0 +1,95 @@
+"""Single migration: create ContestRole, copy contest authors/curators/testers data, remove old M2M fields."""
+import django.db.models.deletion
+from django.db import migrations, models
+
+ROLE_AUTHOR = 'A'
+ROLE_CURATOR = 'C'
+ROLE_TESTER = 'T'
+
+
+def forwards(apps, schema_editor):
+    Contest = apps.get_model('judge', 'Contest')
+    ContestRole = apps.get_model('judge', 'ContestRole')
+
+    contest_roles = []
+    for contest in Contest.objects.prefetch_related('authors', 'curators', 'testers').iterator():
+        for author in contest.authors.all():
+            contest_roles.append(ContestRole(contest=contest, user=author, role=ROLE_AUTHOR))
+        for curator in contest.curators.all():
+            contest_roles.append(ContestRole(contest=contest, user=curator, role=ROLE_CURATOR))
+        for tester in contest.testers.all():
+            contest_roles.append(ContestRole(contest=contest, user=tester, role=ROLE_TESTER))
+    ContestRole.objects.bulk_create(contest_roles, ignore_conflicts=True)
+
+
+def backwards(apps, schema_editor):
+    Contest = apps.get_model('judge', 'Contest')
+    ContestRole = apps.get_model('judge', 'ContestRole')
+
+    for contest in Contest.objects.all().iterator():
+        roles = ContestRole.objects.filter(contest=contest)
+        contest.authors.set(roles.filter(role=ROLE_AUTHOR).values_list('user', flat=True))
+        contest.curators.set(roles.filter(role=ROLE_CURATOR).values_list('user', flat=True))
+        contest.testers.set(roles.filter(role=ROLE_TESTER).values_list('user', flat=True))
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ('judge', '0219_problemdata_zipfile_size_alter_contest_authors_and_more'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='ContestRole',
+            fields=[
+                (
+                    'id',
+                    models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID'),
+                ),
+                (
+                    'role',
+                    models.CharField(
+                        choices=[('A', 'Author'), ('C', 'Curator'), ('T', 'Tester')],
+                        max_length=1,
+                        verbose_name='role',
+                    ),
+                ),
+                (
+                    'contest',
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE,
+                        related_name='contest_roles',
+                        to='judge.contest',
+                        verbose_name='contest',
+                    ),
+                ),
+                (
+                    'user',
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE,
+                        related_name='contest_roles',
+                        to='judge.profile',
+                        verbose_name='user',
+                    ),
+                ),
+            ],
+            options={
+                'verbose_name': 'contest role',
+                'verbose_name_plural': 'contest roles',
+                'unique_together': {('contest', 'user', 'role')},
+            },
+        ),
+        migrations.RunPython(forwards, backwards),
+        migrations.RemoveField(
+            model_name='contest',
+            name='authors',
+        ),
+        migrations.RemoveField(
+            model_name='contest',
+            name='curators',
+        ),
+        migrations.RemoveField(
+            model_name='contest',
+            name='testers',
+        ),
+    ]

judge/models/__init__.py

@@ -11,6 +11,7 @@
     problem_directory_file
 from judge.models.profile import Badge, Organization, OrganizationMonthlyUsage, OrganizationRequest, \
     Profile, WebAuthnCredential
+from judge.models.role import ContestRole, ROLE_AUTHOR, ROLE_CURATOR, ROLE_TESTER
 from judge.models.runtime import Judge, Language, RuntimeVersion
 from judge.models.submission import SUBMISSION_RESULT, Submission, SubmissionSource, SubmissionTestCase
 from judge.models.tag import Tag, TagData, TagGroup, TagProblem
@@ -21,6 +22,7 @@
 revisions.register(LanguageLimit)
 revisions.register(Contest, follow=['contest_problems'])
 revisions.register(ContestProblem)
+revisions.register(ContestRole)
 revisions.register(Organization)
 revisions.register(BlogPost)
 revisions.register(Solution)

judge/models/contest.py

@@ -18,12 +18,13 @@
 from judge import contest_format, event_poster as event
 from judge.models.problem import Problem
 from judge.models.profile import Organization, Profile
+from judge.models.role import ContestRole, RoleQuerySetAdapter, ROLE_AUTHOR, ROLE_CURATOR, ROLE_TESTER
 from judge.models.submission import Submission
 from judge.ratings import rate_contest
 from judge.utils.unicode import utf8bytes
 
 __all__ = ['Contest', 'ContestTag', 'ContestAnnouncement', 'ContestParticipation', 'ContestProblem',
-           'ContestSubmission', 'Rating']
+           'ContestSubmission', 'Rating', 'ContestRole']
 
 
 class MinValueOrNoneValidator(MinValueValidator):
@@ -74,14 +75,6 @@ class Contest(models.Model):
     key = models.CharField(max_length=32, verbose_name=_('contest id'), unique=True,
                            validators=[RegexValidator('^[a-z0-9_]+$', _('Contest id must be ^[a-z0-9_]+$'))])
     name = models.CharField(max_length=100, verbose_name=_('contest name'), db_index=True)
-    authors = models.ManyToManyField(Profile, help_text=_('These users will be able to edit the contest.'),
-                                     related_name='authors+')
-    curators = models.ManyToManyField(Profile, help_text=_('These users will be able to edit the contest, '
-                                                           'but will not be listed as authors.'),
-                                      related_name='curators+', blank=True)
-    testers = models.ManyToManyField(Profile, help_text=_('These users will be able to view the contest, '
-                                                          'but not edit it.'),
-                                     blank=True, related_name='testers+')
     description = models.TextField(verbose_name=_('description'), blank=True)
     problems = models.ManyToManyField(Problem, verbose_name=_('problems'), through='ContestProblem')
     start_time = models.DateTimeField(verbose_name=_('start time'), db_index=True)
@@ -318,6 +311,23 @@ def show_scoreboard(self):
             return False
         return True
 
+    def _role_users(self, role):
+        return RoleQuerySetAdapter(Profile.objects.filter(
+            contest_roles__contest=self, contest_roles__role=role,
+        ))
+
+    @property
+    def authors(self):
+        return self._role_users(ROLE_AUTHOR)
+
+    @property
+    def curators(self):
+        return self._role_users(ROLE_CURATOR)
+
+    @property
+    def testers(self):
+        return self._role_users(ROLE_TESTER)
+
     @property
     def contest_window_length(self):
         return self.end_time - self.start_time
@@ -375,18 +385,22 @@ def time_before_end(self):
     def ended(self):
         return self.end_time < self._now
 
+    def _role_ids(self, roles):
+        return ContestRole.objects.filter(
+            contest=self, role__in=roles,
+        ).values_list('user_id', flat=True)
+
     @cached_property
     def author_ids(self):
-        return Contest.authors.through.objects.filter(contest=self).values_list('profile_id', flat=True)
+        return self._role_ids([ROLE_AUTHOR])
 
     @cached_property
     def editor_ids(self):
-        return self.author_ids.union(
-            Contest.curators.through.objects.filter(contest=self).values_list('profile_id', flat=True))
+        return self._role_ids([ROLE_AUTHOR, ROLE_CURATOR])
 
     @cached_property
     def tester_ids(self):
-        return Contest.testers.through.objects.filter(contest=self).values_list('profile_id', flat=True)
+        return self._role_ids([ROLE_TESTER])
 
     @classmethod
     def get_id_secret(cls, contest_id):
@@ -530,31 +544,9 @@ def get_visible_contests(cls, user):
                 )
             )
 
-            authors_exists = Contest.authors.through.objects.filter(
-                contest_id=OuterRef('pk'),
-                profile_id=user.profile.id,
-            )
-            curators_exists = Contest.curators.through.objects.filter(
-                contest_id=OuterRef('pk'),
-                profile_id=user.profile.id,
-            )
-            testers_exists = Contest.testers.through.objects.filter(
-                contest_id=OuterRef('pk'),
-                profile_id=user.profile.id,
-            )
-
             queryset = queryset.annotate(
-                has_author=Exists(authors_exists),
-                has_curator=Exists(curators_exists),
-                has_tester=Exists(testers_exists),
-            )
-
-            queryset = queryset.filter(
-                q |
-                Q(has_author=True) |
-                Q(has_curator=True) |
-                Q(has_tester=True),
-            )
+                has_role=ContestRole.exists_for(user.profile),
+            ).filter(q | Q(has_role=True))
 
         return queryset
 

judge/models/role.py

@@ -0,0 +1,69 @@
+from django.db import models
+from django.db.models import CASCADE, Exists, OuterRef
+from django.utils.translation import gettext_lazy as _
+
+from judge.models.profile import Profile
+
+ROLE_AUTHOR = 'A'
+ROLE_CURATOR = 'C'
+ROLE_TESTER = 'T'
+
+ROLE_CHOICES = (
+    (ROLE_AUTHOR, _('Author')),
+    (ROLE_CURATOR, _('Curator')),
+    (ROLE_TESTER, _('Tester')),
+)
+
+
+class RoleQuerySetAdapter:
+    def __init__(self, queryset):
+        self.queryset = queryset
+
+    def all(self):
+        return self.queryset.all()
+
+    def filter(self, *args, **kwargs):
+        return self.queryset.filter(*args, **kwargs)
+
+    def exists(self):
+        return self.queryset.exists()
+
+    def first(self):
+        return self.queryset.first()
+
+    def values_list(self, *args, **kwargs):
+        return self.queryset.values_list(*args, **kwargs)
+
+    def __iter__(self):
+        return iter(self.queryset)
+
+    def __contains__(self, item):
+        return self.queryset.filter(pk=getattr(item, 'pk', item)).exists()
+
+
+class ContestRole(models.Model):
+    contest = models.ForeignKey(
+        'Contest', verbose_name=_('contest'), related_name='contest_roles', on_delete=CASCADE,
+    )
+    user = models.ForeignKey(
+        Profile, verbose_name=_('user'), related_name='contest_roles', on_delete=CASCADE,
+    )
+    role = models.CharField(max_length=1, choices=ROLE_CHOICES, verbose_name=_('role'))
+
+    class Meta:
+        unique_together = ('contest', 'user', 'role')
+        verbose_name = _('contest role')
+        verbose_name_plural = _('contest roles')
+
+    def __str__(self):
+        return f'{self.user} - {self.get_role_display()} of {self.contest}'
+
+    @staticmethod
+    def exists_for(user, role=None, roles=None):
+        filters = {'contest_id': OuterRef('pk'), 'user': user}
+        if role:
+            filters['role'] = role
+        qs = ContestRole.objects.filter(**filters)
+        if roles:
+            qs = qs.filter(role__in=roles)
+        return Exists(qs)