diff --git a/src/backend/core/migrations/0027_item_quota_excluded.py b/src/backend/core/migrations/0027_item_quota_excluded.py new file mode 100644 index 00000000..88de409f --- /dev/null +++ b/src/backend/core/migrations/0027_item_quota_excluded.py @@ -0,0 +1,16 @@ +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('core', '0026_item_item_creator_size_not_hdel_idx'), + ] + + operations = [ + migrations.AddField( + model_name='item', + name='quota_excluded', + field=models.BooleanField(default=False, help_text="Exclude this item from its creator's storage quota computation."), + ), + ] diff --git a/src/backend/core/migrations/0028_item_creator_size_quota_idx.py b/src/backend/core/migrations/0028_item_creator_size_quota_idx.py new file mode 100644 index 00000000..143c9406 --- /dev/null +++ b/src/backend/core/migrations/0028_item_creator_size_quota_idx.py @@ -0,0 +1,26 @@ +from django.contrib.postgres.operations import AddIndexConcurrently, RemoveIndexConcurrently +from django.db import migrations, models + + +class Migration(migrations.Migration): + + # CREATE/DROP INDEX CONCURRENTLY cannot run inside a transaction. It avoids + # locking writes on the item table while the index is being built. The new + # index is added before the old one is removed so the storage used + # aggregation never loses index support during the deployment. + atomic = False + + dependencies = [ + ('core', '0027_item_quota_excluded'), + ] + + operations = [ + AddIndexConcurrently( + model_name='item', + index=models.Index(condition=models.Q(('hard_deleted_at__isnull', True), ('quota_excluded', False)), fields=['creator'], include=('size',), name='item_creator_size_quota_idx'), + ), + RemoveIndexConcurrently( + model_name='item', + name='item_creator_size_not_hdel_idx', + ), + ] diff --git a/src/backend/core/models.py b/src/backend/core/models.py index 993f8502..aacdf51d 100644 --- a/src/backend/core/models.py +++ b/src/backend/core/models.py @@ -49,7 +49,7 @@ from wopi.conversion.policy import target_extension_for logger = getLogger(__name__) # Item fields whose update can change the storage used by a user. -STORAGE_USED_FIELDS = {"size", "creator", "creator_id", "hard_deleted_at"} +STORAGE_USED_FIELDS = {"size", "creator", "creator_id", "hard_deleted_at", "quota_excluded"} def get_trashbin_cutoff(): @@ -1021,6 +1021,10 @@ class Item(TreeModel, BaseModel): mimetype = models.CharField(max_length=255, null=True, blank=True) main_workspace = models.BooleanField(default=False) size = models.BigIntegerField(null=True, blank=True) + quota_excluded = models.BooleanField( + default=False, + help_text=_("Exclude this item from its creator's storage quota computation."), + ) description = models.TextField(null=True, blank=True) malware_detection_info = models.JSONField( null=True, @@ -1054,8 +1058,8 @@ class Item(TreeModel, BaseModel): models.Index( fields=["creator"], include=["size"], - condition=models.Q(hard_deleted_at__isnull=True), - name="item_creator_size_not_hdel_idx", + condition=models.Q(hard_deleted_at__isnull=True, quota_excluded=False), + name="item_creator_size_quota_idx", ), ] diff --git a/src/backend/core/storage/creator_storage_compute_backend.py b/src/backend/core/storage/creator_storage_compute_backend.py index c2fb7082..d2fcfe04 100644 --- a/src/backend/core/storage/creator_storage_compute_backend.py +++ b/src/backend/core/storage/creator_storage_compute_backend.py @@ -15,6 +15,6 @@ class CreatorStorageComputeBackend(StorageComputeBackend): """ Compute the total storage used by a set of users. """ - return Item.objects.filter(creator__in=users, hard_deleted_at__isnull=True).aggregate( - total_size=Sum("size", default=0) - )["total_size"] + return Item.objects.filter( + creator__in=users, hard_deleted_at__isnull=True, quota_excluded=False + ).aggregate(total_size=Sum("size", default=0))["total_size"] diff --git a/src/backend/core/tests/storage/test_creator_storage_compute_backend.py b/src/backend/core/tests/storage/test_creator_storage_compute_backend.py index 21d9b6f0..9bacc3dd 100644 --- a/src/backend/core/tests/storage/test_creator_storage_compute_backend.py +++ b/src/backend/core/tests/storage/test_creator_storage_compute_backend.py @@ -39,3 +39,21 @@ def test_compute_storage_used_keeps_soft_deleted_items(): soft_deleted.soft_delete() assert CreatorStorageComputeBackend().compute_storage_used([user]) == 350 + + +def test_compute_storage_used_excludes_quota_excluded_items(): + """Items flagged as quota excluded should not count toward the storage used.""" + user = factories.UserFactory() + factories.ItemFactory(creator=user, size=100) + factories.ItemFactory(creator=user, size=250, quota_excluded=True) + + assert CreatorStorageComputeBackend().compute_storage_used([user]) == 100 + + +def test_compute_storage_used_only_quota_excluded_items(): + """A user owning only quota excluded items should have a storage used of zero.""" + user = factories.UserFactory() + factories.ItemFactory(creator=user, size=100, quota_excluded=True) + factories.ItemFactory(creator=user, size=250, quota_excluded=True) + + assert CreatorStorageComputeBackend().compute_storage_used([user]) == 0 diff --git a/src/backend/core/tests/test_api_entitlements_local.py b/src/backend/core/tests/test_api_entitlements_local.py index 379d0f8b..ba66f24e 100644 --- a/src/backend/core/tests/test_api_entitlements_local.py +++ b/src/backend/core/tests/test_api_entitlements_local.py @@ -419,6 +419,31 @@ def test_api_entitlements_local_cache_invalidated_on_item_save( assert response.json()["quota"]["usage"] == 700 +@override_settings( + ENTITLEMENTS_BACKEND=LOCAL_BACKEND, + ENTITLEMENTS_BACKEND_PARAMETERS={"default_storage_limit": 1000}, +) +def test_api_entitlements_local_cache_invalidated_on_quota_excluded( + django_capture_on_commit_callbacks, +): + """Flagging an item as quota excluded should invalidate the creator's cached usage.""" + client = APIClient() + user = factories.UserFactory() + item = factories.ItemFactory(type=models.ItemTypeChoices.FILE, creator=user, size=400) + client.force_authenticate(user) + + response = client.get("/api/v1.0/entitlements/") + assert response.json()["quota"]["usage"] == 400 + + item.quota_excluded = True + with django_capture_on_commit_callbacks(execute=True): + item.save(update_fields=["quota_excluded"]) + + assert cache.get(get_storage_used_cache_key(user.id)) is None + response = client.get("/api/v1.0/entitlements/") + assert response.json()["quota"]["usage"] == 0 + + def test_api_entitlements_local_cache_invalidated_on_hard_delete( django_capture_on_commit_callbacks, ):