models: Move currently_used_upload_space function to Realm model.

zerver/lib/upload.py

@@ -4,7 +4,6 @@ from django.utils.translation import ugettext as _
 from django.conf import settings
 from django.core.files import File
 from django.http import HttpRequest
-from django.db.models import Sum
 from jinja2 import Markup as mark_safe
 import unicodedata
 
@@ -274,17 +273,11 @@ def upload_image_to_s3(
 
     key.set_contents_from_string(contents, headers=headers)  # type: ignore # https://github.com/python/typeshed/issues/1552
 
-def currently_used_upload_space(realm: Realm) -> int:
-    used_space = Attachment.objects.filter(realm=realm).aggregate(Sum('size'))['size__sum']
-    if used_space is None:
-        return 0
-    return used_space
-
 def check_upload_within_quota(realm: Realm, uploaded_file_size: int) -> None:
     upload_quota = realm.upload_quota_bytes()
     if upload_quota is None:
         return
-    used_space = currently_used_upload_space(realm)
+    used_space = realm.currently_used_upload_space_bytes()
     if (used_space + uploaded_file_size) > upload_quota:
         raise RealmUploadQuotaError(_("Upload would exceed your organization's upload quota."))
 

zerver/models.py

@@ -4,7 +4,7 @@ from typing.re import Match
 
 from django.db import models
 from django.db.models.query import QuerySet
-from django.db.models import Manager, CASCADE
+from django.db.models import Manager, Sum, CASCADE
 from django.conf import settings
 from django.contrib.auth.models import AbstractBaseUser, UserManager, \
     PermissionsMixin
@@ -394,6 +394,12 @@ class Realm(models.Model):
         # it as gibibytes (GiB) to be a bit more generous in case of confusion.
         return self.upload_quota_gb << 30
 
+    def currently_used_upload_space_bytes(self) -> int:
+        used_space = Attachment.objects.filter(realm=self).aggregate(Sum('size'))['size__sum']
+        if used_space is None:
+            return 0
+        return used_space
+
     @property
     def subdomain(self) -> str:
         return self.string_id

zerver/tests/test_upload.py

@@ -23,8 +23,7 @@ from zerver.lib.upload import sanitize_name, S3UploadBackend, \
     upload_message_file, upload_emoji_image, delete_message_image, LocalUploadBackend, \
     ZulipUploadBackend, MEDIUM_AVATAR_SIZE, resize_avatar, \
     resize_emoji, BadImageError, get_realm_for_filename, \
-    currently_used_upload_space, DEFAULT_AVATAR_SIZE, DEFAULT_EMOJI_SIZE, \
-    exif_rotate
+    DEFAULT_AVATAR_SIZE, DEFAULT_EMOJI_SIZE, exif_rotate
 import zerver.lib.upload
 from zerver.models import Attachment, get_user, \
     Message, UserProfile, Realm, \
@@ -1780,15 +1779,15 @@ class UploadSpaceTests(UploadSerializeMixin, ZulipTestCase):
         self.user_profile = self.example_user('hamlet')
 
     def test_currently_used_upload_space(self) -> None:
-        self.assertEqual(0, currently_used_upload_space(self.realm))
+        self.assertEqual(0, self.realm.currently_used_upload_space_bytes())
 
         data = b'zulip!'
         upload_message_file(u'dummy.txt', len(data), u'text/plain', data, self.user_profile)
-        self.assertEqual(len(data), currently_used_upload_space(self.realm))
+        self.assertEqual(len(data), self.realm.currently_used_upload_space_bytes())
 
         data2 = b'more-data!'
         upload_message_file(u'dummy2.txt', len(data2), u'text/plain', data2, self.user_profile)
-        self.assertEqual(len(data) + len(data2), currently_used_upload_space(self.realm))
+        self.assertEqual(len(data) + len(data2), self.realm.currently_used_upload_space_bytes())
 
 class ExifRotateTests(TestCase):
     def test_image_do_not_rotate(self) -> None: