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

@@ -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:

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:

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())