From f7346f36fcbf21b256bd615f7558c22efa2dfafc Mon Sep 17 00:00:00 2001 From: Ujjawal Modi Date: Fri, 9 Jun 2023 01:16:38 +0530 Subject: [PATCH] attachments: Refactor code for flushing used_upload_space cache. Subsequent commits will add "on_delete=models.RESTRICT" relationships, which will result in the Attachment objects being deleted after Realm has been deleted from the database. In order to handle this, we update get_realm_used_upload_space_cache_key function to accept realm_id as parameter instead of realm object, so that the code for flushing the cache works even after the realm is deleted. This change is fine because eventually only realm_id is used by this function and there is no need of the complete realm object. --- zerver/lib/cache.py | 6 +++--- zerver/models.py | 4 +++- zerver/tests/test_upload.py | 14 ++++++++------ 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/zerver/lib/cache.py b/zerver/lib/cache.py index d16be41a8b..d2916ad02d 100644 --- a/zerver/lib/cache.py +++ b/zerver/lib/cache.py @@ -498,8 +498,8 @@ def get_muting_users_cache_key(muted_user_id: int) -> str: return f"muting_users_list:{muted_user_id}" -def get_realm_used_upload_space_cache_key(realm: "Realm") -> str: - return f"realm_used_upload_space:{realm.id}" +def get_realm_used_upload_space_cache_key(realm_id: int) -> str: + return f"realm_used_upload_space:{realm_id}" def active_user_ids_cache_key(realm_id: int) -> str: @@ -699,7 +699,7 @@ def flush_used_upload_space_cache( attachment = instance if created: - cache_delete(get_realm_used_upload_space_cache_key(attachment.owner.realm)) + cache_delete(get_realm_used_upload_space_cache_key(attachment.owner.realm_id)) def to_dict_cache_key_id(message_id: int) -> str: diff --git a/zerver/models.py b/zerver/models.py index 164eb6db19..bd37e4967a 100644 --- a/zerver/models.py +++ b/zerver/models.py @@ -950,7 +950,9 @@ class Realm(models.Model): # type: ignore[django-manager-missing] # django-stub # `realm` instead of `self` here to make sure the parameters of the cache key # function matches the original method. - @cache_with_key(get_realm_used_upload_space_cache_key, timeout=3600 * 24 * 7) + @cache_with_key( + lambda realm: get_realm_used_upload_space_cache_key(realm.id), timeout=3600 * 24 * 7 + ) def currently_used_upload_space_bytes(realm) -> int: # noqa: N805 used_space = Attachment.objects.filter(realm=realm).aggregate(Sum("size"))["size__sum"] if used_space is None: diff --git a/zerver/tests/test_upload.py b/zerver/tests/test_upload.py index 6ee54dd6a9..a780cc0a5c 100644 --- a/zerver/tests/test_upload.py +++ b/zerver/tests/test_upload.py @@ -1865,21 +1865,22 @@ 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(None, cache_get(get_realm_used_upload_space_cache_key(self.realm.id))) 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]) + self.assertEqual(0, cache_get(get_realm_used_upload_space_cache_key(self.realm.id))[0]) data = b"zulip!" upload_message_attachment("dummy.txt", len(data), "text/plain", data, self.user_profile) # notify_attachment_update function calls currently_used_upload_space_bytes which # updates the cache. - self.assert_length(data, cache_get(get_realm_used_upload_space_cache_key(self.realm))[0]) + self.assert_length(data, cache_get(get_realm_used_upload_space_cache_key(self.realm.id))[0]) self.assert_length(data, self.realm.currently_used_upload_space_bytes()) data2 = b"more-data!" upload_message_attachment("dummy2.txt", len(data2), "text/plain", data2, self.user_profile) self.assertEqual( - len(data) + len(data2), cache_get(get_realm_used_upload_space_cache_key(self.realm))[0] + len(data) + len(data2), + cache_get(get_realm_used_upload_space_cache_key(self.realm.id))[0], ) self.assertEqual(len(data) + len(data2), self.realm.currently_used_upload_space_bytes()) @@ -1887,12 +1888,13 @@ class UploadSpaceTests(UploadSerializeMixin, ZulipTestCase): 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] + len(data) + len(data2), + cache_get(get_realm_used_upload_space_cache_key(self.realm.id))[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(None, cache_get(get_realm_used_upload_space_cache_key(self.realm.id))) self.assert_length(data2, self.realm.currently_used_upload_space_bytes())