[ENG-8186] Storage allocation for draft registrations

admin/draft_registrations/urls.py

@@ -9,4 +9,6 @@
 urlpatterns = [
     re_path(r'^$', views.UserDraftRegistrationSearchView.as_view(), name='search'),
     re_path(r'^(?P<draft_registration_id>\w+)/$', views.DraftRegistrationView.as_view(), name='detail'),
+    re_path(r'^(?P<draft_registration_id>\w+)/modify_storage_usage/$', views.DraftRegisrationModifyStorageUsage.as_view(),
+        name='adjust-draft-registration-storage-usage'),
 ]

admin/draft_registrations/views.py

@@ -6,13 +6,20 @@
 from django.views.generic import DetailView
 
 from admin.base.forms import GuidForm
+from admin.nodes.queries import STORAGE_USAGE_QUERY
+from admin.nodes.views import StorageMixin
+
 from osf.models.registrations import DraftRegistration
 
 
 class DraftRegistrationMixin(PermissionRequiredMixin):
 
     def get_object(self):
-        draft_registration = DraftRegistration.load(self.kwargs['draft_registration_id'])
+        draft_registration = DraftRegistration.objects.filter(
+            _id=self.kwargs['draft_registration_id']
+        ).annotate(
+            **STORAGE_USAGE_QUERY
+        ).first()
         draft_registration.guid = draft_registration._id
         return draft_registration
 
@@ -52,3 +59,8 @@ def get_context_data(self, **kwargs):
         return super().get_context_data(**{
             'draft_registration': draft_registration
         }, **kwargs)
+
+
+class DraftRegisrationModifyStorageUsage(DraftRegistrationMixin, StorageMixin):
+    template_name = 'draft_registrations/detail.html'
+    permission_required = 'osf.change_draftregistration'

admin/nodes/queries.py

@@ -0,0 +1,29 @@
+from django.db.models import F, Case, When, IntegerField
+
+from website import settings
+
+
+STORAGE_USAGE_QUERY = {
+    'public_cap': Case(
+        When(
+            custom_storage_usage_limit_public=None,
+            then=settings.STORAGE_LIMIT_PUBLIC,
+        ),
+        When(
+            custom_storage_usage_limit_public__gt=0,
+            then=F('custom_storage_usage_limit_public'),
+        ),
+        output_field=IntegerField()
+    ),
+    'private_cap': Case(
+        When(
+            custom_storage_usage_limit_private=None,
+            then=settings.STORAGE_LIMIT_PRIVATE,
+        ),
+        When(
+            custom_storage_usage_limit_private__gt=0,
+            then=F('custom_storage_usage_limit_private'),
+        ),
+        output_field=IntegerField()
+    )
+}

admin/nodes/views.py

@@ -6,7 +6,7 @@
 from django.utils import timezone
 from django.core.exceptions import PermissionDenied, ValidationError
 from django.urls import NoReverseMatch
-from django.db.models import F, Case, When, IntegerField
+from django.db.models import F
 from django.contrib import messages
 from django.contrib.auth.mixins import PermissionRequiredMixin
 from django.http import HttpResponse
@@ -23,6 +23,7 @@
 from admin.base.views import GuidView
 from admin.base.forms import GuidForm
 from admin.notifications.views import detect_duplicate_notifications, delete_selected_notifications
+from admin.nodes.queries import STORAGE_USAGE_QUERY
 
 from api.share.utils import update_share
 from api.caching.tasks import update_storage_usage_cache
@@ -61,34 +62,33 @@ def get_object(self):
             guids___id=self.kwargs['guid']
         ).annotate(
             guid=F('guids___id'),
-            public_cap=Case(
-                When(
-                    custom_storage_usage_limit_public=None,
-                    then=settings.STORAGE_LIMIT_PUBLIC,
-                ),
-                When(
-                    custom_storage_usage_limit_public__gt=0,
-                    then=F('custom_storage_usage_limit_public'),
-                ),
-                output_field=IntegerField()
-            ),
-            private_cap=Case(
-                When(
-                    custom_storage_usage_limit_private=None,
-                    then=settings.STORAGE_LIMIT_PRIVATE,
-                ),
-                When(
-                    custom_storage_usage_limit_private__gt=0,
-                    then=F('custom_storage_usage_limit_private'),
-                ),
-                output_field=IntegerField()
-            )
+            **STORAGE_USAGE_QUERY
         ).get()
 
     def get_success_url(self):
         return reverse('nodes:node', kwargs={'guid': self.kwargs['guid']})
 
 
+class StorageMixin(View):
+
+    def post(self, request, *args, **kwargs):
+        object = self.get_object()
+        new_private_cap = request.POST.get('private-cap-input')
+        new_public_cap = request.POST.get('public-cap-input')
+
+        object_private_cap = object.custom_storage_usage_limit_private or settings.STORAGE_LIMIT_PRIVATE
+        object_public_cap = object.custom_storage_usage_limit_public or settings.STORAGE_LIMIT_PUBLIC
+
+        if float(new_private_cap) != object_private_cap:
+            object.custom_storage_usage_limit_private = new_private_cap
+
+        if float(new_public_cap) != object_public_cap:
+            object.custom_storage_usage_limit_public = new_public_cap
+
+        object.save()
+        return redirect(self.get_success_url())
+
+
 class NodeView(NodeMixin, GuidView):
     """ Allows authorized users to view node info.
     """
@@ -601,28 +601,11 @@ def post(self, request, *args, **kwargs):
         return redirect(self.get_success_url())
 
 
-class NodeModifyStorageUsage(NodeMixin, View):
+class NodeModifyStorageUsage(NodeMixin, StorageMixin):
     """ Allows an authorized user to view a node's storage usage info and set their public/private storage cap.
     """
     permission_required = 'osf.change_node'
 
-    def post(self, request, *args, **kwargs):
-        node = self.get_object()
-        new_private_cap = request.POST.get('private-cap-input')
-        new_public_cap = request.POST.get('public-cap-input')
-
-        node_private_cap = node.custom_storage_usage_limit_private or settings.STORAGE_LIMIT_PRIVATE
-        node_public_cap = node.custom_storage_usage_limit_public or settings.STORAGE_LIMIT_PUBLIC
-
-        if float(new_private_cap) != node_private_cap:
-            node.custom_storage_usage_limit_private = new_private_cap
-
-        if float(new_public_cap) != node_public_cap:
-            node.custom_storage_usage_limit_public = new_public_cap
-
-        node.save()
-        return redirect(self.get_success_url())
-
 
 class NodeRecalculateStorage(NodeMixin, View):
     """ Allows an authorized user to manually set a node's storage cache by recalculating the value.

admin/templates/draft_registrations/detail.html

@@ -54,12 +54,7 @@ <h2>Draft Registration: <b>{{ draft_registration.title }}</b> <a href="{{ draft_
                         {% endif %}
                     </tr>
                     {% include "draft_registrations/contributors.html" with draft_registration=draft_registration %}
-                    <tr>
-                        <td>Node storage usage</td>
-                        <td>
-                            <b>Current usage:</b> {{ draft_registration.storage_usage }}<br>
-                        </td>
-                    </tr>
+                    {% include "draft_registrations/storage_usage.html" with draft_registration=draft_registration %}
                 </tbody>
             </table>
         </div>

admin/templates/draft_registrations/storage_usage.html

@@ -0,0 +1,35 @@
+{% load node_extras %}
+
+<tr>
+    <td>Draft registration storage usage</td>
+    <td>
+        <b>Public cap:</b> {{ draft_registration.public_cap|floatformat:0 }} GB<br>
+        <b>Private cap:</b> {{ draft_registration.private_cap|floatformat:0 }} GB<br>
+        <a href="{% url 'draft_registrations:adjust-draft-registration-storage-usage' draft_registration_id=draft_registration.guid %}"
+           data-toggle="modal" data-target="#modifyStorageCaps"
+           class="btn btn-warning">
+            Modify Storage Caps
+        </a>
+        <div class="modal" id="modifyStorageCaps">
+            <div class="modal-dialog">
+                <div class="modal-content">
+                    <form class="well" method="post" action="{% url 'draft_registrations:adjust-draft-registration-storage-usage' draft_registration_id=draft_registration.guid %}">
+                    {% csrf_token %}
+                    <div class="modal-header">
+                        <button type="button" class="close" data-dismiss="modal">x</button>
+                        <h3>Adjusting the storage caps for {{ draft_registration.guid }}</h3>
+                    </div>
+                    <b>Public cap:</b> <input name='public-cap-input' type="text" value="{{ draft_registration.public_cap|floatformat:0 }}" /> GB<br>
+                    <b>Private cap: </b><input name='private-cap-input' type="text" value="{{ draft_registration.private_cap|floatformat:0 }}" /> GB
+                    <div class="modal-footer">
+                        <input class="btn btn-success" type="submit" value="Save" />
+                        <button type="button" class="btn btn-default" data-dismiss="modal">
+                            Cancel
+                        </button>
+                    </div>
+                </form>
+                </div>
+            </div>
+        </div>
+    </td>
+</tr>

osf/migrations/0030_draftregistration_custom_storage_usage_limit_private_and_more.py

@@ -0,0 +1,23 @@
+# Generated by Django 4.2.15 on 2025-06-17 13:42
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('osf', '0029_remove_abstractnode_keenio_read_key'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='draftregistration',
+            name='custom_storage_usage_limit_private',
+            field=models.DecimalField(blank=True, decimal_places=9, max_digits=100, null=True),
+        ),
+        migrations.AddField(
+            model_name='draftregistration',
+            name='custom_storage_usage_limit_public',
+            field=models.DecimalField(blank=True, decimal_places=9, max_digits=100, null=True),
+        ),
+    ]

osf/models/archive.py

@@ -146,7 +146,7 @@ def _set_target(self, addon_short_name):
 
     def set_targets(self):
         addons = []
-        for addon in [self.src_node.get_addon(name)
+        for addon in [self.src_node.get_addon(name, cached=False)
                       for name in settings.ADDONS_ARCHIVABLE
                       if settings.ADDONS_ARCHIVABLE[name] != 'none']:
             if not addon or not isinstance(addon, BaseStorageAddon) or not addon.complete:

osf/models/mixins.py

@@ -554,12 +554,12 @@ def get_or_add_addon(self, name, *args, **kwargs):
             return addon
         return self.add_addon(name, *args, **kwargs)
 
-    def get_addon(self, name, is_deleted=False, auth=None):
+    def get_addon(self, name, is_deleted=False, auth=None, cached=True):
         # Avoid test-breakages by avoiding early access to the request context
         if name not in self.OSF_HOSTED_ADDONS:
             request, user_id = get_request_and_user_id()
             if flag_is_active(request, features.ENABLE_GV):
-                return self._get_addon_from_gv(gv_pk=name, requesting_user_id=user_id, auth=auth)
+                return self._get_addon_from_gv(gv_pk=name, requesting_user_id=user_id, auth=auth, cached=cached)
 
         try:
             settings_model = self._settings_model(name)

osf/models/node.py

@@ -2466,21 +2466,28 @@ def _remove_from_associated_collections(self, auth=None, force=False):
                     force=True
                 )
 
-    def _get_addon_from_gv(self, gv_pk, requesting_user_id, auth=None):
+    def _get_addons_from_gv_without_caching(self, gv_pk, requesting_user_id, auth=None):
+        requesting_user = OSFUser.load(requesting_user_id)
+        services = gv_translations.get_external_services(requesting_user)
+        for service in services:
+            if service.short_name == gv_pk:
+                break
+        else:
+            return None
+
+        return self._get_addons_from_gv(requesting_user_id, service.type, auth=auth)
+
+    def _get_addon_from_gv(self, gv_pk, requesting_user_id, auth=None, cached=True):
         request = get_current_request()
         # This is to avoid making multiple requests to GV
         # within the lifespan of one request on the OSF side
-        try:
-            gv_addons = request.gv_addons
-        except AttributeError:
-            requesting_user = OSFUser.load(requesting_user_id)
-            services = gv_translations.get_external_services(requesting_user)
-            for service in services:
-                if service.short_name == gv_pk:
-                    break
-            else:
-                return None
-            gv_addons = request.gv_addons = self._get_addons_from_gv(requesting_user_id, service.type, auth=auth)
+        if cached:
+            try:
+                gv_addons = request.gv_addons
+            except AttributeError:
+                gv_addons = request.gv_addons = self._get_addons_from_gv_without_caching(gv_pk, requesting_user_id, auth=auth)
+        else:
+            gv_addons = self._get_addons_from_gv_without_caching(gv_pk, requesting_user_id, auth=auth)
 
         for item in gv_addons:
             if item.short_name == gv_pk:

osf/models/registrations.py

@@ -980,6 +980,9 @@ class DraftRegistration(ObjectIDMixin, RegistrationResponseMixin, DirtyFieldsMix
         default=get_default_id,
     )
 
+    custom_storage_usage_limit_public = models.DecimalField(decimal_places=9, max_digits=100, null=True, blank=True)
+    custom_storage_usage_limit_private = models.DecimalField(decimal_places=9, max_digits=100, null=True, blank=True)
+
     # Dictionary field mapping question id to a question's comments and answer
     # {
     #   <qid>: {