Add user-based quota enforcement

dkc/core/admin/__init__.py

@@ -1,5 +1,6 @@
 from .file import FileAdmin
 from .folder import FolderAdmin
+from .quota import QuotaAdmin
 from .terms import TermsAdmin
 
-__all__ = ['FileAdmin', 'FolderAdmin', 'TermsAdmin']
+__all__ = ['FileAdmin', 'FolderAdmin', 'QuotaAdmin', 'TermsAdmin']

dkc/core/admin/quota.py

@@ -0,0 +1,39 @@
+from django.contrib import admin
+from django.db import models
+from django_admin_display import admin_display
+import humanize
+
+from dkc.core.models import Quota
+
+
+@admin.register(Quota)
+class QuotaAdmin(admin.ModelAdmin):
+    list_display = ['id', 'user', 'human_used', 'human_allowed', 'usage_percent']
+    list_select_related = ['user']
+
+    search_fields = ['user__username']
+
+    readonly_fields = ['used', 'usage_percent']
+
+    @admin_display(short_description='Used', admin_order_field='used')
+    def human_used(self, obj: Quota) -> str:
+        return humanize.naturalsize(obj.used, binary=True)
+
+    @admin_display(short_description='Allowed', admin_order_field='used')
+    def human_allowed(self, obj: Quota) -> str:
+        return humanize.naturalsize(obj.allowed, binary=True)
+
+    @admin_display(
+        short_description='% Used',
+        admin_order_field=models.Case(
+            # Prevent division by zero, and sort the result as effectively 100%
+            models.When(allowed=0, then=1),
+            # Multiply by 1.0 to force floating-point division
+            default=(models.F('used') * 1.0 / models.F('allowed')),
+            output_field=models.FloatField(),
+        ),
+    )
+    def usage_percent(self, obj: Quota) -> str:
+        if obj.allowed == 0:
+            return '--'
+        return '{:.1%}'.format(obj.used / obj.allowed)

dkc/core/migrations/0008_quota.py

@@ -0,0 +1,62 @@
+# Generated by Django 3.1.5 on 2021-01-31 23:46
+
+from django.conf import settings
+from django.db import migrations, models
+import django.db.models.deletion
+import django.db.models.expressions
+
+import dkc.core.models.quota
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ('core', '0007_non_nullable_creator'),
+    ]
+
+    operations = [
+        migrations.RemoveField(
+            model_name='quota',
+            name='allocation',
+        ),
+        migrations.AddField(
+            model_name='quota',
+            name='allowed',
+            field=models.PositiveBigIntegerField(default=dkc.core.models.quota._default_user_quota),
+        ),
+        migrations.AddField(
+            model_name='quota',
+            name='used',
+            field=models.PositiveBigIntegerField(default=0),
+        ),
+        migrations.AddField(
+            model_name='tree',
+            name='quota',
+            field=models.ForeignKey(
+                default=1,
+                on_delete=django.db.models.deletion.PROTECT,
+                related_name='trees',
+                to='core.quota',
+            ),
+            preserve_default=False,
+        ),
+        migrations.AlterField(
+            model_name='quota',
+            name='user',
+            field=models.OneToOneField(
+                default=1,
+                on_delete=django.db.models.deletion.CASCADE,
+                related_name='quota',
+                to='auth.user',
+            ),
+            preserve_default=False,
+        ),
+        migrations.AddConstraint(
+            model_name='quota',
+            constraint=models.CheckConstraint(
+                check=models.Q(used__lte=django.db.models.expressions.F('allowed')),
+                name='used_lte_allowed',
+            ),
+        ),
+    ]

dkc/core/models/folder.py

@@ -5,7 +5,7 @@
 from django.contrib.auth.models import User
 from django.core import validators
 from django.core.exceptions import ValidationError
-from django.db import models
+from django.db import models, transaction
 from django.dispatch import receiver
 from django_extensions.db.models import TimeStampedModel
 from girder_utils.models import JSONObjectField
@@ -101,16 +101,20 @@ def abs_path(self) -> str:
     def public(self) -> bool:
         return self.tree.public
 
+    @transaction.atomic
     def increment_size(self, amount: int) -> None:
         """
         Increment or decrement the Folder size.
 
-        This Folder and all of its ancestors' sizes will be updated atomically. ``amount`` may be
-        negative, but an operation resulting in a negative final size is illegal.
+        This Folder, all of its ancestors' sizes, and its quota will be updated atomically.
+        ``amount`` may be negative, but an operation resulting in a negative final size is illegal.
         """
         if amount == 0:
             return
 
+        # Do this first, in case it fails
+        self.tree.quota.increment(amount)
+
         ancestor_pks = [folder.pk for folder in self.ancestors]
         Folder.objects.filter(pk__in=ancestor_pks).update(size=(models.F('size') + amount))
 

dkc/core/models/quota.py

@@ -1,26 +1,62 @@
 from typing import Type
 
+from django.conf import settings
 from django.contrib.auth.models import User
-from django.db import models
+from django.core.exceptions import ValidationError
+from django.db import IntegrityError, models, transaction
 from django.db.models.signals import post_save
 from django.dispatch import receiver
-from girder_utils.models import SelectRelatedManager
+
+
+def _default_user_quota() -> int:
+    return settings.DKC_DEFAULT_QUOTA
 
 
 class Quota(models.Model):
-    allocation = models.BigIntegerField(default=0)
-    user = models.OneToOneField(User, null=True, on_delete=models.CASCADE)
-    # root_set
+    class Meta:
+        constraints = [
+            models.CheckConstraint(
+                check=models.Q(used__lte=models.F('allowed')),
+                name='used_lte_allowed',
+            )
+        ]
+
+    used = models.PositiveBigIntegerField(default=0)
+    allowed = models.PositiveBigIntegerField(default=_default_user_quota)
+
+    # OneToOneField ensures that only one Quota per User exists
+    user = models.OneToOneField(User, on_delete=models.CASCADE)
+
+    @transaction.atomic
+    def increment(self, amount: int) -> None:
+        """
+        Increment or decrement the quota.
+
+        This function increments (or decrements, if ``amount`` is negative) the usage
+        values associated with this quota. If this increment would exceed the max
+        allotment for this quota, a ``ValidationError`` is raised and the transaction
+        is rolled back.
+        """
+        if amount == 0:
+            return
 
-    # The user is needed to stringify, so join it when directly querying for a quota
-    objects = SelectRelatedManager('user')
+        try:
+            # Use an .update query instead of a .save, to avoid assigning an F-expression on the
+            # local instance, which might need to be rolled back on a failure
+            Quota.objects.filter(pk=self.pk).update(used=(models.F('used') + amount))
+        except IntegrityError as e:
+            if '"used_lte_allowed"' in str(e):
+                raise ValidationError(
+                    'Root folder size quota would be exceeded: ' f'{self.used}B > {self.allowed}B.'
+                )
+            else:
+                raise
 
-    def __str__(self):
-        return f'Quota for <{self.user}>; {self.allocation} bytes'
+        # Update with the new value
+        self.refresh_from_db(fields=['used'])
 
 
 @receiver(post_save, sender=User)
-def create_user_quota(sender: Type[User], instance: User, created: bool, **kwargs):
+def _create_user_quota(sender: Type[User], instance: User, created: bool, **kwargs):
     if created:
-        quota = Quota(user=instance)
-        quota.save()
+        Quota.objects.create(user=instance)

dkc/core/models/tree.py

@@ -20,13 +20,17 @@
 
 from dkc.core.permissions import Permission, PermissionGrant
 
+from .quota import Quota
+
 if TYPE_CHECKING:
     # Prevent circular import
     from .folder import Folder
 
 
 class Tree(models.Model):
     public: bool = models.BooleanField(default=False)
+    # Prevent deletion of a Quota if it has Trees referencing it
+    quota = models.ForeignKey(Quota, on_delete=models.PROTECT, related_name='trees')
 
     @property
     def root_folder(self) -> Folder:

dkc/core/rest/folder.py

@@ -182,7 +182,8 @@ def perform_create(self, serializer: FolderSerializer):
             if not tree.has_permission(user, permission=Permission.write):
                 raise PermissionDenied()
         else:
-            tree = Tree.objects.create()
+            # Use the user's quota for the new tree
+            tree = Tree.objects.create(quota=user.quota)
             tree.grant_permission(PermissionGrant(user_or_group=user, permission=Permission.admin))
         serializer.save(tree=tree, creator=user)
 

dkc/core/tests/factories.py

@@ -20,10 +20,15 @@ class Meta:
 
 class TreeFactory(factory.django.DjangoModelFactory):
     public = False
+    # No need to instantiate a quota, just fetch from the user creating this
+    quota = factory.SelfAttribute('creator.quota')
 
     class Meta:
         model = Tree
 
+    class Params:
+        creator = factory.SubFactory(UserFactory)
+
 
 class FolderFactory(factory.django.DjangoModelFactory):
     class Meta:
@@ -34,7 +39,10 @@ class Meta:
     user_metadata = _metadata_faker
     parent = None
     tree = factory.Maybe(
-        'parent', factory.SelfAttribute('parent.tree'), factory.SubFactory(TreeFactory)
+        'parent',
+        factory.SelfAttribute('parent.tree'),
+        # Make the new tree be created by (and use the quota of) this folder's creator
+        factory.SubFactory(TreeFactory, creator=factory.SelfAttribute('..creator')),
     )
     creator = factory.SubFactory(UserFactory)
 
@@ -51,13 +59,11 @@ class Meta:
     creator = factory.SubFactory(UserFactory)
 
 
-class TreeWithRootFactory(factory.django.DjangoModelFactory):
-    class Meta:
-        model = Tree
-
+class TreeWithRootFactory(TreeFactory):
     @factory.post_generation
     def root_folder(self, create, *args, **kwargs):
-        return FolderFactory(tree=self)
+        # Make the new folder be created by the same user as it's tree's quota
+        return FolderFactory(tree=self, creator=self.quota.user)
 
 
 class TermsFactory(factory.django.DjangoModelFactory):

dkc/core/tests/test_file.py

@@ -10,9 +10,8 @@ def test_file_abs_path(folder, folder_factory, file_factory):
     assert grandchild.abs_path == f'/{folder.name}/{child.name}/{grandchild.name}'
 
 
-def test_file_checksum(file_factory):
-    # Use "build" strategy, so database is not required
-    file = file_factory.build()
+@pytest.mark.django_db
+def test_file_checksum(file):
     file.compute_sha512()
     assert len(file.sha512) == 128
 

dkc/core/tests/test_folder.py

@@ -5,24 +5,22 @@
 from dkc.core.models.folder import MAX_DEPTH
 
 
+@pytest.mark.django_db
 def test_folder_name_invalid(folder_factory):
-    folder = folder_factory.build(name='name / withslash')
+    folder = folder_factory(name='name / withslash')
 
-    # Since the folder is not saved and added to a tree, other validation errors are also present,
-    # so it's critical to match the error by string content
     with pytest.raises(ValidationError, match='Name may not contain forward slashes'):
         folder.full_clean()
 
 
-def test_is_root_root(folder_factory):
-    folder = folder_factory.build()
+@pytest.mark.django_db
+def test_is_root_root(folder):
     assert folder.is_root is True
 
 
-def test_is_root_child(folder_factory):
-    folder = folder_factory.build()
-    child = folder_factory.build(parent=folder)
-    assert child.is_root is False
+@pytest.mark.django_db
+def test_is_root_child(child_folder):
+    assert child_folder.is_root is False
 
 
 @pytest.mark.django_db
@@ -87,11 +85,9 @@ def test_folder_sibling_names_unique_files(file, folder_factory):
 
 @pytest.mark.django_db
 def test_root_folder_names_unique(folder, folder_factory):
-    other_root = folder_factory.build(name=folder.name)
-    # Make sure foreign key references exist in database first
-    other_root.creator.save()
-    other_root.tree.save()
-    with pytest.raises(IntegrityError):
+    other_root = folder_factory()
+    other_root.name = folder.name
+    with pytest.raises(IntegrityError, match=r'folder_name'):
         other_root.save()
 
 
@@ -107,6 +103,8 @@ def test_folder_names_not_globally_unique(folder_factory):
 def test_increment_size(folder_factory, amount):
     initial_size = 100
     root = folder_factory(size=initial_size)
+    root.tree.quota.used = initial_size
+    root.tree.quota.save()
     child = folder_factory(parent=root, size=initial_size)
     grandchild = folder_factory(parent=child, size=initial_size)
 
@@ -118,14 +116,18 @@ def test_increment_size(folder_factory, amount):
     assert grandchild.size == new_size
     assert grandchild.parent.size == new_size
     assert grandchild.parent.parent.size == new_size
+    assert grandchild.parent.parent.tree.quota.used == new_size
 
 
 @pytest.mark.django_db
 def test_increment_size_negative(folder_factory):
     # Make the root too small
     root = folder_factory(size=5)
+    root.tree.quota.used = 10
+    root.tree.quota.save()
     child = folder_factory(parent=root, size=10)
 
     # Increment the child, which tests enforcement across propagation
-    with pytest.raises(IntegrityError, match=r'size'):
+    # An IntegrityError is expected, since this should cause a 500 if it actually happens
+    with pytest.raises(IntegrityError, match=r'folder_size'):
         child.increment_size(-10)