Merge pull request #5601 from ozer550/fix-perseus-storage-calculation

Remove perseus files from storage calculations

Author: bjester

contentcuration/contentcuration/models.py

@@ -353,25 +353,77 @@ 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")
-            .distinct(),
-            name="files",
+
+        user_files_cte = With(
+            self.files.get_queryset().values(
+                "id",
+                "checksum",
+                "contentnode_id",
+                "file_format_id",
+                "file_size",
+                "preset_id",
+            ),
+            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"))
+                )
+            )
+        )
+
+        existing_checksums_cte = With(
+            editable_files_qs.values("checksum").distinct(),
+            name="existing_checksums",
+        )
+
+        staging_files_qs = (
+            user_files_cte.queryset()
+            .with_cte(user_files_cte)
+            .filter(
+                Exists(
+                    ContentNode.objects.filter(
+                        tree_id=channel.staging_tree.tree_id,
+                        id=OuterRef("contentnode_id"),
+                    )
+                )
+            )
+        )
+
+        new_staging_files_qs = (
+            staging_files_qs.with_cte(tree_cte)
+            .with_cte(existing_checksums_cte)
+            .exclude(
+                Exists(
+                    existing_checksums_cte.queryset().filter(
+                        checksum=OuterRef("checksum"),
+                    )
+                )
+            )
+        )
+
+        new_staging_files_qs = self._filter_storage_billable_files(new_staging_files_qs)
+
+        unique_staging_ids = (
+            new_staging_files_qs.order_by("checksum", "id")
+            .distinct("checksum")
+            .values("id")
         )
         staged_size = float(
-            staging_tree_files.aggregate(used=Sum("file_size"))["used"] or 0
+            new_staging_files_qs.filter(id__in=Subquery(unique_staging_ids)).aggregate(
+                used=Sum("file_size")
+            )["used"]
+            or 0
         )
 
         if self.get_available_space() < staged_size:
@@ -414,13 +466,55 @@ 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().only(
+                "id",
+                "checksum",
+                "contentnode_id",
+                "file_format_id",
+                "file_size",
+                "preset_id",
+            ),
+            name="user_files",
+        )
+
+        base_files_qs = (
+            user_files_cte.queryset()
+            .with_cte(tree_cte)
+            .with_cte(user_files_cte)
+            .filter(
+                Exists(
+                    tree_cte.join(
+                        ContentNode.objects.only("id", "tree_id"),
+                        tree_id=tree_cte.col.tree_id,
+                    )
+                    .with_cte(tree_cte)
+                    .filter(id=OuterRef("contentnode_id"))
+                )
+            )
+        )
+
+        base_files_qs = self._filter_storage_billable_files(base_files_qs)
+
+        unique_file_ids = (
+            base_files_qs.order_by("checksum", "id").distinct("checksum").values("id")
+        )
+
+        files_qs = base_files_qs.filter(id__in=Subquery(unique_file_ids))
+
+        return files_qs
+
+    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):

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)