Improve query that checks storage space before channel deploy

contentcuration/contentcuration/models.py

@@ -347,22 +347,37 @@ def check_feature_flag(self, flag_name):
         return feature_flags.get(flag_name, False)
 
     def check_channel_space(self, channel):
-        active_files = self.get_user_active_files()
-        staging_tree_id = channel.staging_tree.tree_id
-        channel_files = (
-            self.files.filter(contentnode__tree_id=staging_tree_id)
+        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",
+        )
+
+        staging_tree_files = (
+            self.files.filter(contentnode__tree_id=channel.staging_tree.tree_id)
+            .with_cte(tree_cte)
+            .with_cte(files_cte)
+            .exclude(Exists(files_cte.queryset().filter(checksum=OuterRef("checksum"))))
             .values("checksum")
             .distinct()
-            .exclude(checksum__in=active_files.values_list("checksum", flat=True))
         )
-        staged_size = float(channel_files.aggregate(used=Sum("file_size"))["used"] or 0)
+        staged_size = float(
+            staging_tree_files.aggregate(used=Sum("file_size"))["used"] or 0
+        )
 
-        if self.get_available_space(active_files=active_files) < (staged_size):
+        if self.get_available_space() < staged_size:
             raise PermissionDenied(
                 _("Out of storage! Request more space under Settings > Storage.")
             )
 
     def check_staged_space(self, size, checksum):
+        """
+        .. deprecated:: only used in `api_file_upload` which is now deprecated
+        """
         if self.staged_files.filter(checksum=checksum).exists():
             return True
         space = self.get_available_staged_space()
@@ -372,6 +387,9 @@ def check_staged_space(self, size, checksum):
             )
 
     def get_available_staged_space(self):
+        """
+        .. deprecated:: only used in `api_file_upload` which is now deprecated
+        """
         space_used = (
             self.staged_files.values("checksum")
             .distinct()

contentcuration/contentcuration/ricecooker_versions.py

@@ -3,15 +3,15 @@
 Any version >= VERSION_OK will get a message that
 the version is "up to date" (log level = info)
 """
-VERSION_OK = "0.6.32"  # this gets overwritten to current v. after XML RPC call
+VERSION_OK = "0.7.3"
 VERSION_OK_MESSAGE = "Ricecooker v{} is up-to-date."
 
 """
 Recommended minimum ricecooker version
 Any version < VERSION_OK and >= VERSION_SOFT_WARNING will get a
 recommendation to upgrade before running (log level = warning)
 """
-VERSION_SOFT_WARNING = "0.6.30"
+VERSION_SOFT_WARNING = "0.7.0"
 VERSION_SOFT_WARNING_MESSAGE = (
     "You are using Ricecooker v{}, however v{} is available. "
     "You should consider upgrading your Ricecooker."

contentcuration/contentcuration/tests/test_files.py

@@ -1,20 +1,29 @@
 # -*- coding: utf-8 -*-
 import json
+from uuid import uuid4
 
+import mock
 import pytest
+from django.core.exceptions import PermissionDenied
 from django.core.files.storage import default_storage
+from django.db.models import Exists
+from django.db.models import OuterRef
 from le_utils.constants import content_kinds
 from mock import patch
 
 from .base import BaseAPITestCase
 from .base import StudioTestCase
 from .testdata import base64encoding
 from .testdata import generated_base64encoding
+from .testdata import node
 from contentcuration.api import write_raw_content_to_storage
+from contentcuration.models import Channel
 from contentcuration.models import ContentNode
 from contentcuration.models import delete_empty_file_reference
 from contentcuration.models import File
 from contentcuration.models import generate_object_storage_name
+from contentcuration.models import StagedFile
+from contentcuration.models import User
 from contentcuration.utils.files import create_thumbnail_from_base64
 from contentcuration.utils.files import get_thumbnail_encoding
 from contentcuration.utils.nodes import map_files_to_node
@@ -82,3 +91,144 @@ def test_delete_empty_file_reference(self):
         assert default_storage.exists(storage_path), "file should be saved"
         delete_empty_file_reference(checksum, "pdf")
         assert not default_storage.exists(storage_path), "file should be deleted"
+
+
+class StagedChannelSpaceTestCase(StudioTestCase):
+    """
+    Tests for
+    - User.check_channel_space()
+    - User.get_available_staged_space()
+    - User.check_staged_space()
+    """
+
+    def setUp(self):
+        super().setUpBase()
+
+        self.staged_channel = Channel.objects.create(
+            name="Staged", actor_id=self.user.id, language_id="en"
+        )
+        self.staged_channel.save()
+
+        file_node_id = uuid4().hex
+        self.staged_channel.staging_tree = node(
+            {
+                "node_id": uuid4().hex,
+                "kind_id": "topic",
+                "title": "Root Node",
+                "children": [
+                    {
+                        "node_id": file_node_id,
+                        "kind_id": "video",
+                        "title": "Video 1",
+                    }
+                ],
+            },
+            parent=None,
+        )
+        self.staged_channel.save()
+        self.node = ContentNode.objects.get(node_id=file_node_id)
+        self._set_uploader(self.channel)
+        self._set_uploader(
+            self.staged_channel, self.staged_channel.staging_tree.tree_id
+        )
+        self.node_file = self.node.files.all()[0]
+
+    def _set_uploader(self, channel: Channel, tree_id=None):
+        if tree_id is None:
+            tree_id = channel.main_tree.tree_id
+
+        File.objects.filter(
+            Exists(
+                ContentNode.objects.filter(
+                    tree_id=tree_id, id=OuterRef("contentnode_id")
+                )
+            )
+        ).update(uploaded_by=self.user)
+
+    def _create_duplicate(self, file: File):
+        dupe_node = node(
+            {
+                "node_id": uuid4().hex,
+                "kind_id": "video",
+                "title": "Video 2",
+            },
+            parent=self.node.parent,
+        )
+        dupe_file = dupe_node.files.all()[0]
+        dupe_file.file_size = file.file_size
+        dupe_file.checksum = file.checksum
+        dupe_file.uploaded_by = self.user
+        dupe_file.save(set_by_file_on_disk=False)
+
+    def test_check_channel_space__okay(self):
+        try:
+            self.user.check_channel_space(self.staged_channel)
+        except PermissionDenied:
+            self.fail("Staging channel space is larger than available")
+
+    def test_check_channel_space__duplicate_checksum_same_tree(self):
+        # set file to slightly more than half, such that if both files are included, it should
+        # exceed the available space
+        self.node_file.file_size = self.user.disk_space / 2 + 1
+        self.node_file.checksum = uuid4().hex
+        self.node_file.save(set_by_file_on_disk=False)
+        self._create_duplicate(self.node_file)
+
+        try:
+            self.user.check_channel_space(self.staged_channel)
+        except PermissionDenied:
+            self.fail("Staging channel space is larger than available")
+
+    def test_check_channel_space__duplicate_checksum_different_tree(self):
+        # set file larger than space
+        self.node_file.file_size = self.user.disk_space + 1
+        self.node_file.save(set_by_file_on_disk=False)
+
+        # ensure file has matching checksum to another file in deployed channel tree,
+        # which should be the case because of how the test fixtures function
+        deployed_file_count = File.objects.filter(
+            Exists(
+                ContentNode.objects.filter(
+                    tree_id=self.channel.main_tree.tree_id,
+                    id=OuterRef("contentnode_id"),
+                )
+            ),
+            checksum=self.node_file.checksum,
+        ).count()
+        self.assertGreaterEqual(deployed_file_count, 1)
+
+        try:
+            self.user.check_channel_space(self.staged_channel)
+        except PermissionDenied:
+            self.fail("Staging channel space is larger than available")
+
+    def test_check_channel_space__fail(self):
+        self.node_file.file_size = self.user.disk_space + 1
+        self.node_file.checksum = uuid4().hex
+        self.node_file.save(set_by_file_on_disk=False)
+
+        with self.assertRaises(PermissionDenied):
+            self.user.check_channel_space(self.staged_channel)
+
+    def test_get_available_staged_space(self):
+        f = StagedFile.objects.create(
+            checksum=uuid4().hex,
+            uploaded_by=self.user,
+            file_size=100,
+        )
+        expected_available_space = self.user.disk_space - f.file_size
+        self.assertEqual(
+            expected_available_space, self.user.get_available_staged_space()
+        )
+
+    def test_check_staged_space__exists(self):
+        f = StagedFile.objects.create(
+            checksum=uuid4().hex,
+            uploaded_by=self.user,
+            file_size=100,
+        )
+        with mock.patch.object(
+            User, "get_available_staged_space"
+        ) as get_available_staged_space:
+            get_available_staged_space.return_value = 0
+            self.assertTrue(self.user.check_staged_space(100, f.checksum))

contentcuration/contentcuration/views/internal.py

@@ -169,7 +169,10 @@ def file_diff(request):
 @authentication_classes((TokenAuthentication,))
 @permission_classes((IsAuthenticated,))
 def api_file_upload(request):
-    """ Upload a file to the storage system """
+    """Upload a file to the storage system
+
+    .. deprecated:: Ricecooker 0.7+ no longer uses this endpoint
+    """
     try:
         fobj = request.FILES["file"]
         checksum, ext = fobj._name.split(".")