[GRDM-57690] Update user Quota value when moving/copying files

addons/base/views.py

@@ -433,6 +433,9 @@ def get_auth(auth, **kwargs):
         credentials = node.serialize_waterbutler_credentials(provider_name)
         waterbutler_settings = node.serialize_waterbutler_settings(provider_name)
 
+    if is_node_process and provider_settings:
+        waterbutler_settings['max_file_size'] = getattr(provider_settings.config, 'max_file_size', None)
+
     if not is_node_process:
         # for only location_id value
         storage = ExportDataLocation.objects.get(pk=location_id)

addons/osfstorage/decorators.py

@@ -105,7 +105,8 @@ def wrapped(payload, *args, **kwargs):
                 'source': source,
                 'destination': dest_parent,
                 'name': payload['destination']['name'],
-                'is_check_permission': is_check_permission
+                'is_check_permission': is_check_permission,
+                'replaced_size': int(payload.get('replaced_size', 0)),
             })
         except KeyError:
             raise HTTPError(http_status.HTTP_400_BAD_REQUEST)

addons/osfstorage/views.py

@@ -1,6 +1,6 @@
 from __future__ import unicode_literals
 
-from django.db.models import IntegerField
+from django.db.models import IntegerField, Sum
 from django.db.models.functions import Cast
 from rest_framework import status as http_status
 import logging
@@ -19,7 +19,7 @@
 
 from api.caching.tasks import update_storage_usage
 from osf.exceptions import InvalidTagError, TagNotFoundError
-from osf.models import FileVersion, OSFUser, ExportDataRestore, ExportData
+from osf.models import FileVersion, OSFUser, ExportDataRestore, ExportData, FileInfo, BaseFileNode
 from osf.utils.permissions import WRITE
 from osf.utils.requests import check_select_for_update
 from website.project.decorators import (
@@ -28,6 +28,7 @@
 from website.project.model import has_anonymous_link
 
 from website.files import exceptions
+from website.util.quota import get_project_storage_type, update_quota
 from addons.osfstorage import utils
 from addons.osfstorage import decorators
 from addons.osfstorage.models import OsfStorageFolder
@@ -124,17 +125,94 @@ def osfstorage_get_revisions(file_node, payload, target, **kwargs):
         ]
     }
 
+def _get_all_descendant_file_ids(folder_node):
+    TRASHED_TYPES = ('osf.trashedfilenode', 'osf.trashedfile', 'osf.trashedfolder')
+
+    file_ids = []
+    folder_ids_to_process = [folder_node.id]
+
+    while folder_ids_to_process:
+        children = BaseFileNode.objects.filter(
+            parent_id__in=folder_ids_to_process,
+            provider='osfstorage',
+        ).exclude(
+            type__in=TRASHED_TYPES
+        ).values('id', 'type')
+
+        next_folder_ids = []
+        for child in children:
+            if child['type'] == 'osf.osfstoragefile':
+                file_ids.append(child['id'])
+            elif child['type'] == 'osf.osfstoragefolder':
+                next_folder_ids.append(child['id'])
+
+        folder_ids_to_process = next_folder_ids
+
+    return file_ids
 
 @decorators.waterbutler_opt_hook
 def osfstorage_copy_hook(source, destination, name=None, **kwargs):
-    ret = source.copy_under(destination, name=name).serialize(), http_status.HTTP_201_CREATED
+    replaced_size = int(kwargs.get('replaced_size', 0))
+    cloned = source.copy_under(destination, name=name)
+
+    # Only handle if source is a node and destination is a node
+    if getattr(source.target, 'type', '') == 'osf.node' and getattr(destination.target, 'type', '') == 'osf.node':
+        if source.is_file:
+            latest_version = cloned.versions.order_by('-created').first()
+            new_size = latest_version.size if latest_version and latest_version.size else 0
+            FileInfo.objects.update_or_create(file=cloned, defaults={'file_size': new_size})
+        else:
+            source_file_ids = _get_all_descendant_file_ids(source)
+            source_size_map = {
+                fi.file_id: fi.file_size
+                for fi in FileInfo.objects.filter(file_id__in=source_file_ids)
+            }
+
+            cloned_file_ids = _get_all_descendant_file_ids(cloned)
+            cloned_files = list(BaseFileNode.objects.filter(
+                id__in=cloned_file_ids
+            ).values('id', 'copied_from_id'))
+
+            file_infos = []
+            new_size = 0
+            for cf in cloned_files:
+                src_id = cf['copied_from_id']
+                if src_id in source_size_map:
+                    size = source_size_map[src_id]
+                elif src_id is not None:
+                    # FileInfo not exist
+                    node = BaseFileNode.objects.filter(id=src_id).first()
+                    size = 0
+                    if node:
+                        v = node.versions.order_by('-created').first()
+                        size = (v.size or 0) if v else 0
+                else:          # src_id is None and not in map
+                    size = 0
+
+                file_infos.append(FileInfo(file_id=cf['id'], file_size=size))
+                new_size += size
+
+            try:
+                FileInfo.objects.bulk_create(file_infos)
+            except IntegrityError:
+                for fi in file_infos:
+                    FileInfo.objects.get_or_create(file_id=fi.file_id, defaults={'file_size': fi.file_size})
+            replaced_size = 0  # osfstorage_delete already subtracted this
+
+        # UserQuota: delta = new_size - replaced_size
+        delta = new_size - replaced_size
+        if delta != 0:
+            storage_type = get_project_storage_type(destination.target)
+            update_quota(destination.target, abs(delta), storage_type, add=(delta > 0))
+
     update_storage_usage(destination.target)
-    return ret
+    return cloned.serialize(), http_status.HTTP_201_CREATED
 
 @decorators.waterbutler_opt_hook
 def osfstorage_move_hook(source, destination, name=None, **kwargs):
     source_target = source.target
     is_check_permission = kwargs.get('is_check_permission')
+    replaced_size = int(kwargs.get('replaced_size', 0))
     try:
         ret = source.move_under(destination, name=name, is_check_permission=is_check_permission).serialize(), http_status.HTTP_200_OK
     except exceptions.FileNodeCheckedOutError:
@@ -150,6 +228,32 @@ def osfstorage_move_hook(source, destination, name=None, **kwargs):
             'message_long': 'Cannot move file as it is the primary file of preprint.'
         })
 
+    # Only handle if source is a node and destination is a node
+    if getattr(source_target, 'type', '') == 'osf.node' and getattr(destination.target, 'type', '') == 'osf.node':
+        if source.is_file:
+            latest_version = source.versions.order_by('-created').first()
+            new_size = latest_version.size if latest_version and latest_version.size else 0
+            # FileInfo: UPSERT
+            FileInfo.objects.update_or_create(file=source, defaults={'file_size': new_size})
+        else:
+            source_file_ids = _get_all_descendant_file_ids(source)  # recursive
+            new_size = FileInfo.objects.filter(
+                file_id__in=source_file_ids
+            ).aggregate(total=Sum('file_size'))['total'] or 0
+            replaced_size = 0  # osfstorage_delete already subtracted this
+
+        # UserQuota
+        if source_target != destination.target:
+            src_storage = get_project_storage_type(source_target)
+            dest_storage = get_project_storage_type(destination.target)
+            update_quota(source_target, new_size, src_storage, add=False)
+            delta = new_size - replaced_size
+            if delta != 0:
+                update_quota(destination.target, abs(delta), dest_storage, add=(delta > 0))
+        elif replaced_size > 0:
+            storage_type = get_project_storage_type(destination.target)
+            update_quota(destination.target, replaced_size, storage_type, add=False)
+
     # once the move is complete recalculate storage for both targets if it's a inter-target move.
     if is_check_permission and source_target != destination.target:
         update_storage_usage(destination.target)
@@ -379,6 +483,24 @@ def osfstorage_create_child(file_node, payload, **kwargs):
         if not current_version or not current_version.is_duplicate(new_version):
             update_storage_usage(file_node.target)
 
+        # Only handle if target is a node
+        if getattr(file_node.target, 'type', '') == 'osf.node':
+            new_size = new_version.size if new_version.size and new_version.size > 0 else 0
+            if new_size >= 0:
+                old_info = FileInfo.objects.filter(file=file_node).first()
+                old_size = old_info.file_size if old_info else 0
+
+            # FileInfo: UPSERT
+            FileInfo.objects.update_or_create(
+                file=file_node, defaults={'file_size': new_size}
+            )
+
+            # UserQuota: delta = new - old
+            delta = new_size - old_size
+            if delta != 0:
+                storage_type = get_project_storage_type(file_node.target)
+                update_quota(file_node.target, abs(delta), storage_type, add=(delta > 0))
+
         version_id = new_version._id
         archive_exists = new_version.archive is not None
     else:
@@ -406,6 +528,17 @@ def osfstorage_delete(file_node, payload, target, **kwargs):
     if file_node == OsfStorageFolder.objects.get_root(target=target):
             raise HTTPError(http_status.HTTP_400_BAD_REQUEST)
 
+    if getattr(file_node.target, 'type', '') == 'osf.node' and not file_node.is_file:
+        # Since we are deleting a folder, we need to calculate the total size of all descendant files to update quota
+        all_file_ids = _get_all_descendant_file_ids(file_node)
+        total_size = FileInfo.objects.filter(file_id__in=all_file_ids) \
+                                     .aggregate(total=Sum('file_size'))['total'] or 0
+        if total_size > 0:
+            storage_type = get_project_storage_type(target)
+            # Subtract quota BEFORE deletion; also zero FileInfo to avoid race conditions with new file creation during deletion
+            FileInfo.objects.filter(file_id__in=all_file_ids).update(file_size=0)
+            update_quota(target, total_size, storage_type, add=False)
+
     try:
         file_node.delete(user=user)
 

tests/test_addons.py

@@ -113,7 +113,9 @@ def test_auth_download(self):
         data = jwt.decode(jwe.decrypt(res.json['payload'].encode('utf-8'), self.JWE_KEY), settings.WATERBUTLER_JWT_SECRET, algorithm=settings.WATERBUTLER_JWT_ALGORITHM)['data']
         assert_equal(data['auth'], views.make_auth(self.user))
         assert_equal(data['credentials'], self.node_addon.serialize_waterbutler_credentials())
-        assert_equal(data['settings'], self.node_addon.serialize_waterbutler_settings())
+        expected_settings = self.node_addon.serialize_waterbutler_settings().copy()
+        expected_settings['max_file_size'] = 100
+        assert_equal(data['settings'], expected_settings)
         expected_url = furl.furl(self.node.api_url_for('create_waterbutler_log', _absolute=True, _internal=True))
         observed_url = furl.furl(data['callback_url'])
         observed_url.port = expected_url.port
@@ -153,7 +155,9 @@ def test_auth_bad_cookie(self):
         data = jwt.decode(jwe.decrypt(res.json['payload'].encode('utf-8'), self.JWE_KEY), settings.WATERBUTLER_JWT_SECRET, algorithm=settings.WATERBUTLER_JWT_ALGORITHM)['data']
         assert_equal(data['auth'], views.make_auth(self.user))
         assert_equal(data['credentials'], self.node_addon.serialize_waterbutler_credentials())
-        assert_equal(data['settings'], self.node_addon.serialize_waterbutler_settings())
+        expected_settings = self.node_addon.serialize_waterbutler_settings().copy()
+        expected_settings['max_file_size'] = 100
+        assert_equal(data['settings'], expected_settings)
         expected_url = furl.furl(self.node.api_url_for('create_waterbutler_log', _absolute=True, _internal=True))
         observed_url = furl.furl(data['callback_url'])
         observed_url.port = expected_url.port

tests/test_quota.py

@@ -294,9 +294,8 @@ def test_add_file_info(self):
         )
 
         file_info_list = FileInfo.objects.filter(file=self.file).all()
-        assert_equal(file_info_list.count(), 1)
-        file_info = file_info_list.first()
-        assert_equal(file_info.file_size, 1000)
+        # FileInfo was not created
+        assert_equal(file_info_list.count(), 0)
 
     def test_update_file_info(self):
         file_info = FileInfo(file=self.file, file_size=1000)
@@ -414,9 +413,7 @@ def test_add_first_file(self):
             storage_type=UserQuota.NII_STORAGE,
             user=self.project_creator
         ).all()
-        assert_equal(len(user_quota), 1)
-        user_quota = user_quota[0]
-        assert_equal(user_quota.used, 1000)
+        assert_equal(len(user_quota), 0)
 
     def test_add_first_file_custom_storage(self):
         assert_false(UserQuota.objects.filter(user=self.project_creator).exists())
@@ -453,9 +450,7 @@ def test_add_first_file_custom_storage(self):
             storage_type=UserQuota.CUSTOM_STORAGE,
             user=self.project_creator
         ).all()
-        assert_equal(len(user_quota), 1)
-        user_quota = user_quota[0]
-        assert_equal(user_quota.used, 1200)
+        assert_equal(len(user_quota), 0)
 
     def test_add_file(self):
         UserQuota.objects.create(
@@ -492,7 +487,7 @@ def test_add_file(self):
         ).all()
         assert_equal(len(user_quota), 1)
         user_quota = user_quota[0]
-        assert_equal(user_quota.used, 6500)
+        assert_equal(user_quota.used, 5500)
 
     def test_add_file_custom_storage(self):
         UserQuota.objects.create(
@@ -536,7 +531,7 @@ def test_add_file_custom_storage(self):
         ).all()
         assert_equal(len(user_quota), 1)
         user_quota = user_quota[0]
-        assert_equal(user_quota.used, 6700)
+        assert_equal(user_quota.used, 5500)
 
     def test_add_file_negative_size(self):
         quota.update_used_quota(
@@ -1358,20 +1353,22 @@ def test_delete_file_with_Amazon_S3_Compatible_Storage_for_Institution(self):
         if check_select_for_update():
             mock_file_info.objects.filter.return_value.select_for_update.return_value.get.return_value = FileInfo(
                 file=self.base_file_node, file_size=1000)
-            mock_user_quota.objects.filter.return_value.select_for_update.return_value.first.return_value = UserQuota(
+            mock_user_quota.objects.filter.return_value.select_for_update.return_value.get.return_value = UserQuota(
                 user=self.project_creator,
                 storage_type=UserQuota.CUSTOM_STORAGE,
                 max_quota=api_settings.DEFAULT_MAX_QUOTA,
                 used=5500
             )
         else:
             mock_file_info.objects.get.return_value = FileInfo(file=self.base_file_node, file_size=1000)
-            mock_user_quota.objects.filter.return_value.first.return_value = UserQuota(
+            _real_user_quota = UserQuota(
                 user=self.project_creator,
                 storage_type=UserQuota.CUSTOM_STORAGE,
                 max_quota=api_settings.DEFAULT_MAX_QUOTA,
                 used=5500
             )
+            mock_user_quota.objects.get.return_value = _real_user_quota
+            mock_user_quota.objects.filter.return_value.first.return_value = _real_user_quota
         with mock.patch('website.util.quota.BaseFileNode', mock_base_file_node):
             with mock.patch('website.util.quota.FileInfo', mock_file_info):
                 with mock.patch('website.util.quota.UserQuota', mock_user_quota):
@@ -1414,20 +1411,22 @@ def test_delete_folder_with_Amazon_S3_Compatible_Storage_for_Institution(self):
         if check_select_for_update():
             mock_file_info.objects.filter.return_value.select_for_update.return_value.get.return_value = FileInfo(
                 file=self.base_file_node, file_size=1500)
-            mock_user_quota.objects.filter.return_value.select_for_update.return_value.first.return_value = UserQuota(
+            mock_user_quota.objects.filter.return_value.select_for_update.return_value.get.return_value = UserQuota(
                 user=self.project_creator,
                 storage_type=UserQuota.CUSTOM_STORAGE,
                 max_quota=api_settings.DEFAULT_MAX_QUOTA,
                 used=5500
             )
         else:
             mock_file_info.objects.get.return_value = FileInfo(file=self.base_file_node, file_size=1500)
-            mock_user_quota.objects.filter.return_value.first.return_value = UserQuota(
+            _real_user_quota = UserQuota(
                 user=self.project_creator,
                 storage_type=UserQuota.CUSTOM_STORAGE,
                 max_quota=api_settings.DEFAULT_MAX_QUOTA,
                 used=5500
             )
+            mock_user_quota.objects.get.return_value = _real_user_quota
+            mock_user_quota.objects.filter.return_value.first.return_value = _real_user_quota
         with mock.patch('website.util.quota.BaseFileNode', mock_base_file_node):
             with mock.patch('website.util.quota.FileInfo', mock_file_info):
                 with mock.patch('website.util.quota.UserQuota', mock_user_quota):