From 94ae2dc24ebab20b92d5e6000ae7abb76e45fdf9 Mon Sep 17 00:00:00 2001 From: Vishnu Ks Date: Mon, 14 Jan 2019 12:16:31 +0530 Subject: [PATCH] models: Cache currently_used_upload_space_bytes function. --- zerver/lib/cache.py | 9 +++++++++ zerver/models.py | 7 ++++++- zerver/tests/test_upload.py | 18 ++++++++++++++++++ 3 files changed, 33 insertions(+), 1 deletion(-) diff --git a/zerver/lib/cache.py b/zerver/lib/cache.py index 8e0cf076bb..76cef33452 100644 --- a/zerver/lib/cache.py +++ b/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,) diff --git a/zerver/models.py b/zerver/models.py index 626c4e7f48..85f9fa212e 100644 --- a/zerver/models.py +++ b/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) diff --git a/zerver/tests/test_upload.py b/zerver/tests/test_upload.py index 94e8dac714..efa449ae59 100644 --- a/zerver/tests/test_upload.py +++ b/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: