diff --git a/contentcuration/contentcuration/models.py b/contentcuration/contentcuration/models.py index 32727f0159..23adf07846 100644 --- a/contentcuration/contentcuration/models.py +++ b/contentcuration/contentcuration/models.py @@ -349,25 +349,57 @@ def check_feature_flag(self, flag_name): def check_channel_space(self, channel): tree_cte = With(self.get_user_active_trees().distinct(), name="trees") - files_cte = With( - tree_cte.join( - self.files.get_queryset(), contentnode__tree_id=tree_cte.col.tree_id - ) - .values("checksum") + + user_files_cte = With( + self.files.get_queryset() + .values("checksum", "contentnode_id", "file_format_id") .distinct(), - name="files", + name="user_files", ) - staging_tree_files = ( - self.files.filter(contentnode__tree_id=channel.staging_tree.tree_id) + editable_files_qs = ( + user_files_cte.queryset() .with_cte(tree_cte) - .with_cte(files_cte) - .exclude(Exists(files_cte.queryset().filter(checksum=OuterRef("checksum")))) - .values("checksum") - .distinct() + .with_cte(user_files_cte) + .filter( + Exists( + tree_cte.join( + ContentNode.objects.all(), tree_id=tree_cte.col.tree_id + ) + .with_cte(tree_cte) + .filter(id=OuterRef("contentnode_id")) + ) + ) + ) + + editable_files_qs = self._filter_storage_billable_files(editable_files_qs) + + existing_checksums_cte = With( + editable_files_qs.values("checksum").distinct(), name="existing_checksums" ) + + staging_files_qs = self._filter_storage_billable_files( + self.files.filter(contentnode__tree_id=channel.staging_tree.tree_id) + ) + + staging_files_qs = ( + staging_files_qs.with_cte(tree_cte) + .with_cte(user_files_cte) + .with_cte(existing_checksums_cte) + .exclude( + Exists( + existing_checksums_cte.queryset().filter( + checksum=OuterRef("checksum") + ) + ) + ) + ) + staged_size = float( - staging_tree_files.aggregate(used=Sum("file_size"))["used"] or 0 + staging_files_qs.values("checksum") + .distinct() + .aggregate(used=Sum("file_size"))["used"] + or 0 ) if self.get_available_space() < staged_size: @@ -410,13 +442,42 @@ def get_user_active_trees(self): ) def get_user_active_files(self): - cte = With(self.get_user_active_trees().distinct()) - return ( - cte.join(self.files.get_queryset(), contentnode__tree_id=cte.col.tree_id) - .with_cte(cte) - .values("checksum") - .distinct() + tree_cte = With(self.get_user_active_trees().distinct(), name="trees") + + user_files_cte = With( + self.files.get_queryset() + .values("checksum", "contentnode_id", "file_format_id") + .distinct(), + name="user_files", + ) + file_qs = ( + user_files_cte.queryset() + .with_cte(tree_cte) + .with_cte(user_files_cte) + .filter( + Exists( + tree_cte.join( + ContentNode.objects.all(), tree_id=tree_cte.col.tree_id + ) + .with_cte(tree_cte) + .filter(id=OuterRef("contentnode_id")) + ) + ) + ) + + files_qs = self._filter_storage_billable_files(file_qs) + + return files_qs.values("checksum").distinct() + + def _filter_storage_billable_files(self, queryset): + """ + Perseus exports would not be included in storage calculations. + """ + if queryset is None: + return queryset + return queryset.exclude(file_format_id__isnull=True).exclude( + file_format_id=file_formats.PERSEUS ) def get_space_used(self, active_files=None): diff --git a/contentcuration/contentcuration/tests/test_files.py b/contentcuration/contentcuration/tests/test_files.py index 125dcbc48b..7e0b11a573 100755 --- a/contentcuration/contentcuration/tests/test_files.py +++ b/contentcuration/contentcuration/tests/test_files.py @@ -9,6 +9,7 @@ from django.db.models import Exists from django.db.models import OuterRef from le_utils.constants import content_kinds +from le_utils.constants import file_formats from mock import patch from .base import BaseAPITestCase @@ -232,3 +233,57 @@ def test_check_staged_space__exists(self): ) as get_available_staged_space: get_available_staged_space.return_value = 0 self.assertTrue(self.user.check_staged_space(100, f.checksum)) + + def test_check_channel_space_ignores_perseus_exports(self): + with mock.patch("contentcuration.utils.user.calculate_user_storage"): + self.node_file.file_format_id = file_formats.PERSEUS + self.node_file.file_size = self.user.disk_space + 1 + self.node_file.checksum = uuid4().hex + self.node_file.uploaded_by = self.user + self.node_file.save(set_by_file_on_disk=False) + + try: + self.user.check_channel_space(self.staged_channel) + except PermissionDenied: + self.fail("Perseus exports should not count against staging space") + + +class UserStorageUsageTestCase(StudioTestCase): + def setUp(self): + super().setUpBase() + self.contentnode = ( + self.channel.main_tree.get_descendants(include_self=True) + .filter(files__isnull=False) + .first() + ) + self.assertIsNotNone(self.contentnode) + self.base_file = self.contentnode.files.first() + self.assertIsNotNone(self.base_file) + + def _create_file(self, *, file_format, size): + file_record = File( + contentnode=self.contentnode, + checksum=uuid4().hex, + file_format_id=file_format, + file_size=size, + uploaded_by=self.user, + ) + file_record.save(set_by_file_on_disk=False) + return file_record + + def test_get_space_used_excludes_perseus_exports(self): + baseline_usage = self.user.get_space_used() + + perseus_size = 125 + with mock.patch("contentcuration.utils.user.calculate_user_storage"): + self._create_file(file_format=file_formats.PERSEUS, size=perseus_size) + self.assertEqual(self.user.get_space_used(), baseline_usage) + + non_perseus_size = 275 + with mock.patch("contentcuration.utils.user.calculate_user_storage"): + self._create_file( + file_format=self.base_file.file_format_id, size=non_perseus_size + ) + + expected_usage = baseline_usage + non_perseus_size + self.assertEqual(self.user.get_space_used(), expected_usage)