Feat/organization quota limits

dmoj/settings.py

@@ -124,6 +124,10 @@
 VNOJ_MONTHLY_FREE_CREDIT = 3 * 60 * 60
 VNOJ_PRICE_PER_HOUR = 50
 
+# Organization quota limits
+VNOJ_ORGANIZATION_DEFAULT_MAX_PROBLEMS = 1000
+VNOJ_ORGANIZATION_DEFAULT_MAX_STORAGE = 5 * 1024 * 1024 * 1024  # 5GB
+VNOJ_ORGANIZATION_DEFAULT_STORAGE_EXPIRATION_DAYS = 365  # 1 year
 
 VNOJ_LONG_QUEUE_ALERT_THRESHOLD = 10
 

judge/admin/organization.py

@@ -20,7 +20,8 @@ class Meta:
 class OrganizationAdmin(VersionAdmin):
     readonly_fields = ('creation_date', 'current_consumed_credit')
     fields = ('name', 'slug', 'short_name', 'is_open', 'is_unlisted', 'paid_credit', 'current_consumed_credit',
-              'about', 'logo_override_image', 'slots', 'creation_date', 'admins')
+              'about', 'logo_override_image', 'slots', 'max_problems', 'max_storage', 'storage_expiration',
+              'creation_date', 'admins')
     list_display = ('name', 'short_name', 'is_open', 'is_unlisted', 'slots', 'show_public')
     prepopulated_fields = {'slug': ('name',)}
     actions = ('recalculate_points',)

judge/management/commands/backfill_problem_data_size.py

@@ -1,3 +1,6 @@
+import os
+
+from django.conf import settings
 from django.core.management.base import BaseCommand
 
 from judge.models import ProblemData
@@ -19,33 +22,49 @@ def handle(self, *args, **options):
         if dry_run:
             self.stdout.write(self.style.WARNING('DRY RUN MODE - No changes will be saved'))
 
-        problem_data_list = ProblemData.objects.select_related('problem').all()
-        total_count = problem_data_list.count()
+        problem_data_list = ProblemData.objects.all().iterator(chunk_size=5000)
+        total_count = ProblemData.objects.count()
         updated_count = 0
+        batch_size = 1000
+        batch = []
+        root = settings.DMOJ_PROBLEM_DATA_ROOT
 
         self.stdout.write(f'Processing {total_count} problems...\n')
 
         for problem_data in problem_data_list:
-            problem_code = problem_data.problem.code
             old_zipfile_size = problem_data.zipfile_size
+            total_size = 0
 
-            # Calculate new sizes
-            problem_data.update_zipfile_size()
+            # Calculate new sizes directly via OS to bypass Storage abstractions overhead
+            for field in ['zipfile', 'generator', 'custom_checker', 'custom_grader', 'custom_header']:
+                val = getattr(problem_data, field)
+                if val and val.name:
+                    path = os.path.join(root, val.name)
+                    try:
+                        total_size += os.path.getsize(path)
+                    except (OSError, FileNotFoundError):
+                        pass
 
-            new_zipfile_size = problem_data.zipfile_size
+            new_zipfile_size = total_size
 
             # Check if anything changed
             if old_zipfile_size != new_zipfile_size:
+                problem_data.zipfile_size = new_zipfile_size
+                batch.append(problem_data)
                 updated_count += 1
 
                 self.stdout.write(
-                    f'Problem: {problem_code}\n'
-                    f'  Test data: {self._format_size(old_zipfile_size)} -> {self._format_size(new_zipfile_size)}\n',
+                    f'Problem: {problem_data.problem_id}\n'
+                    f'  Total Storage: {self._format_size(old_zipfile_size)} '
+                    f'-> {self._format_size(new_zipfile_size)}\n',
                 )
 
-                if not dry_run:
-                    # Use update_fields to avoid triggering save hooks again
-                    problem_data.save(update_fields=['zipfile_size'])
+                if len(batch) >= batch_size and not dry_run:
+                    ProblemData.objects.bulk_update(batch, ['zipfile_size'])
+                    batch = []
+
+        if batch and not dry_run:
+            ProblemData.objects.bulk_update(batch, ['zipfile_size'])
 
         if dry_run:
             self.stdout.write(self.style.WARNING(f'\nDRY RUN: Would update {updated_count}/{total_count} problems'))

judge/migrations/0220_organization_quota_fields.py

@@ -0,0 +1,33 @@
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('judge', '0219_problemdata_zipfile_size_alter_contest_authors_and_more'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='organization',
+            name='max_problems',
+            field=models.IntegerField(
+                blank=True,
+                default=None,
+                help_text='Maximum number of problems this org can create. Leave blank to use default from settings.',
+                null=True,
+                verbose_name='maximum problems',
+            ),
+        ),
+        migrations.AddField(
+            model_name='organization',
+            name='max_storage',
+            field=models.BigIntegerField(
+                blank=True,
+                default=None,
+                help_text='Maximum storage for test data in bytes. Leave blank to use default from settings.',
+                null=True,
+                verbose_name='maximum storage (bytes)',
+            ),
+        ),
+    ]

judge/migrations/0221_organization_storage_expiration.py

@@ -0,0 +1,22 @@
+# Generated migration for storage_expiration field
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('judge', '0220_organization_quota_fields'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='organization',
+            name='storage_expiration',
+            field=models.DateField(
+                blank=True, default=None,
+                help_text='Expiry date of the paid storage plan. Leave blank for no expiration.',
+                null=True, verbose_name='storage expiration date',
+            ),
+        ),
+    ]

judge/models/problem_data.py

@@ -99,15 +99,15 @@ def has_yml(self):
         return problem_data_storage.exists('%s/init.yml' % self.problem.code)
 
     def update_zipfile_size(self):
-        """Update the zipfile_size field based on the actual file size."""
-        if self.zipfile:
-            try:
-                self.zipfile_size = self.zipfile.size
-            except (OSError, IOError):
-                # If file doesn't exist or can't be accessed, set size to 0
-                self.zipfile_size = 0
-        else:
-            self.zipfile_size = 0
+        """Update the zipfile_size field based on the actual size of all attached files."""
+        total_size = 0
+        for field in [self.zipfile, self.generator, self.custom_checker, self.custom_grader, self.custom_header]:
+            if field:
+                try:
+                    total_size += field.size
+                except (OSError, IOError, ValueError):
+                    pass
+        self.zipfile_size = total_size
 
     def save(self, *args, **kwargs):
         # Update zipfile size before saving

judge/models/profile.py

@@ -76,6 +76,23 @@ class Organization(models.Model):
         default=settings.VNOJ_MONTHLY_FREE_CREDIT,
         help_text=_('Amount of free credits allocated each month'),
     )
+    max_problems = models.IntegerField(
+        default=None, null=True, blank=True,
+        verbose_name=_('maximum problems'),
+        help_text=_('Maximum number of problems this org can create. '
+                    'Leave blank to use default from settings.'),
+    )
+    max_storage = models.BigIntegerField(
+        default=None, null=True, blank=True,
+        verbose_name=_('maximum storage (bytes)'),
+        help_text=_('Maximum storage for test data in bytes. '
+                    'Leave blank to use default from settings.'),
+    )
+    storage_expiration = models.DateField(
+        default=None, null=True, blank=True,
+        verbose_name=_('storage expiration date'),
+        help_text=_('Expiry date of the paid storage plan. Leave blank for no expiration.'),
+    )
 
     _pp_table = [pow(settings.VNOJ_ORG_PP_STEP, i) for i in range(settings.VNOJ_ORG_PP_ENTRIES)]
 
@@ -138,6 +155,50 @@ def consume_credit(self, consumed):
         self.current_consumed_credit += consumed
         self.save(update_fields=['free_credit', 'paid_credit', 'current_consumed_credit'])
 
+    def get_max_problems(self):
+        if self.max_problems is not None:
+            return self.max_problems
+        return settings.VNOJ_ORGANIZATION_DEFAULT_MAX_PROBLEMS
+
+    def get_max_storage(self):
+        if self.is_storage_expired():
+            return settings.VNOJ_ORGANIZATION_DEFAULT_MAX_STORAGE
+        if self.max_storage is not None:
+            return self.max_storage
+        return settings.VNOJ_ORGANIZATION_DEFAULT_MAX_STORAGE
+
+    def is_storage_expired(self):
+        if self.storage_expiration is None:
+            return False
+        return self.storage_expiration < timezone.now().date()
+
+    def get_days_remaining(self):
+        if self.storage_expiration is None:
+            return None
+        remaining = (self.storage_expiration - timezone.now().date()).days
+        return max(remaining, 0)
+
+    def get_current_problem_count(self):
+        return self.problem_set.count()
+
+    def get_current_storage(self):
+        from django.apps import apps
+        ProblemData = apps.get_model('judge', 'ProblemData')
+        result = ProblemData.objects.filter(
+            problem__organization=self,
+        ).aggregate(total=Sum('zipfile_size'))
+        return result['total'] or 0
+
+    def can_create_problem(self):
+        if self.is_storage_expired():
+            return False
+        return self.get_current_problem_count() < self.get_max_problems()
+
+    def can_upload_data(self):
+        if self.is_storage_expired():
+            return False
+        return self.get_current_storage() <= self.get_max_storage()
+
     class Meta:
         ordering = ['name']
         permissions = (