models: Cache currently_used_upload_space_bytes function.

zerver/lib/cache.py

@@ -331,6 +331,9 @@ realm_user_dict_fields = [
 def realm_user_dicts_cache_key(realm_id: int) -> str:
     return "realm_user_dicts:%s" % (realm_id,)
 
+def get_realm_used_upload_space_cache_key(realm: 'Realm') -> str:
+    return u'realm_used_upload_space:%s' % (realm.id,)
+
 def active_user_ids_cache_key(realm_id: int) -> str:
     return "active_user_ids:%s" % (realm_id,)
 
@@ -454,6 +457,12 @@ def flush_stream(sender: Any, **kwargs: Any) -> None:
            Q(default_events_register_stream=stream)).exists():
         cache_delete(bot_dicts_in_realm_cache_key(stream.realm))
 
+def flush_used_upload_space_cache(sender: Any, **kwargs: Any) -> None:
+    attachment = kwargs['instance']
+
+    if kwargs.get("created") is None or kwargs.get("created") is True:
+        cache_delete(get_realm_used_upload_space_cache_key(attachment.owner.realm))
+
 def to_dict_cache_key_id(message_id: int) -> str:
     return 'message_dict:%d' % (message_id,)
 

zerver/models.py

@@ -20,7 +20,8 @@ from zerver.lib.cache import cache_with_key, flush_user_profile, flush_realm, \
     display_recipient_cache_key, cache_delete, active_user_ids_cache_key, \
     get_stream_cache_key, realm_user_dicts_cache_key, \
     bot_dicts_in_realm_cache_key, realm_user_dict_fields, \
-    bot_dict_fields, flush_message, flush_submessage, bot_profile_cache_key
+    bot_dict_fields, flush_message, flush_submessage, bot_profile_cache_key, \
+    flush_used_upload_space_cache, get_realm_used_upload_space_cache_key
 from zerver.lib.utils import make_safe_digest, generate_random_token
 from django.db import transaction
 from django.utils.timezone import now as timezone_now
@@ -394,6 +395,7 @@ class Realm(models.Model):
         # it as gibibytes (GiB) to be a bit more generous in case of confusion.
         return self.upload_quota_gb << 30
 
+    @cache_with_key(get_realm_used_upload_space_cache_key, timeout=3600*24*7)
     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:
@@ -1750,6 +1752,9 @@ class Attachment(AbstractAttachment):
             } for m in self.messages.all()]
         }
 
+post_save.connect(flush_used_upload_space_cache, sender=Attachment)
+post_delete.connect(flush_used_upload_space_cache, sender=Attachment)
+
 def validate_attachment_request(user_profile: UserProfile, path_id: str) -> Optional[bool]:
     try:
         attachment = Attachment.objects.get(path_id=path_id)

zerver/tests/test_upload.py

@@ -34,6 +34,7 @@ from zerver.lib.actions import (
     do_delete_old_unclaimed_attachments,
     internal_send_private_message,
 )
+from zerver.lib.cache import get_realm_used_upload_space_cache_key, cache_get
 from zerver.lib.create_user import copy_user_settings
 from zerver.lib.users import get_api_key
 from zerver.views.upload import upload_file_backend
@@ -1779,15 +1780,32 @@ class UploadSpaceTests(UploadSerializeMixin, ZulipTestCase):
         self.user_profile = self.example_user('hamlet')
 
     def test_currently_used_upload_space(self) -> None:
+        self.assertEqual(None, cache_get(get_realm_used_upload_space_cache_key(self.realm)))
         self.assertEqual(0, self.realm.currently_used_upload_space_bytes())
+        self.assertEqual(0, cache_get(get_realm_used_upload_space_cache_key(self.realm))[0])
 
         data = b'zulip!'
         upload_message_file(u'dummy.txt', len(data), u'text/plain', data, self.user_profile)
+        self.assertEqual(None, cache_get(get_realm_used_upload_space_cache_key(self.realm)))
         self.assertEqual(len(data), self.realm.currently_used_upload_space_bytes())
+        self.assertEqual(len(data), cache_get(get_realm_used_upload_space_cache_key(self.realm))[0])
 
         data2 = b'more-data!'
         upload_message_file(u'dummy2.txt', len(data2), u'text/plain', data2, self.user_profile)
+        self.assertEqual(None, cache_get(get_realm_used_upload_space_cache_key(self.realm)))
         self.assertEqual(len(data) + len(data2), self.realm.currently_used_upload_space_bytes())
+        self.assertEqual(len(data) + len(data2), cache_get(get_realm_used_upload_space_cache_key(self.realm))[0])
+
+        attachment = Attachment.objects.get(file_name="dummy.txt")
+        attachment.file_name = "dummy1.txt"
+        attachment.save(update_fields=["file_name"])
+        self.assertEqual(len(data) + len(data2), cache_get(get_realm_used_upload_space_cache_key(self.realm))[0])
+        self.assertEqual(len(data) + len(data2), self.realm.currently_used_upload_space_bytes())
+
+        attachment.delete()
+        self.assertEqual(None, cache_get(get_realm_used_upload_space_cache_key(self.realm)))
+        self.assertEqual(len(data2), self.realm.currently_used_upload_space_bytes())
+        self.assertEqual(len(data2), cache_get(get_realm_used_upload_space_cache_key(self.realm))[0])
 
 class ExifRotateTests(TestCase):
     def test_image_do_not_rotate(self) -> None: