realm: Change implementation approach for upload_quota_gb.

Most importantly, fixes a bug where a realm with a custom
.upload_quota_gb value (set by changing it in the database via e.g.
manage.py shell) would end up having it lowered while upgrading their
plan via the do_change_realm_plan_type function, which used to just set
it to the value implied by the new plan without caring about whether
that isn't lower than the original limit.

The new approach is cleaner since we don't do db queries by
upload_quota_gb so it's nicer to just generate these dynamically, making
changes to our limit-per-plan rules much easier - skipping the need for
migrations.

zerver/actions/realm_settings.py

@@ -692,23 +692,18 @@ def do_change_realm_plan_type(
     if plan_type == Realm.PLAN_TYPE_PLUS:
         realm.max_invites = Realm.INVITES_STANDARD_REALM_DAILY_MAX
         realm.message_visibility_limit = None
-        realm.upload_quota_gb = Realm.UPLOAD_QUOTA_STANDARD
     elif plan_type == Realm.PLAN_TYPE_STANDARD:
         realm.max_invites = Realm.INVITES_STANDARD_REALM_DAILY_MAX
         realm.message_visibility_limit = None
-        realm.upload_quota_gb = Realm.UPLOAD_QUOTA_STANDARD
     elif plan_type == Realm.PLAN_TYPE_SELF_HOSTED:
         realm.max_invites = None  # type: ignore[assignment] # https://github.com/python/mypy/issues/3004
         realm.message_visibility_limit = None
-        realm.upload_quota_gb = None
     elif plan_type == Realm.PLAN_TYPE_STANDARD_FREE:
         realm.max_invites = Realm.INVITES_STANDARD_REALM_DAILY_MAX
         realm.message_visibility_limit = None
-        realm.upload_quota_gb = Realm.UPLOAD_QUOTA_STANDARD
     elif plan_type == Realm.PLAN_TYPE_LIMITED:
         realm.max_invites = settings.INVITES_DEFAULT_REALM_DAILY_MAX
         realm.message_visibility_limit = Realm.MESSAGE_VISIBILITY_LIMITED
-        realm.upload_quota_gb = Realm.UPLOAD_QUOTA_LIMITED
     else:
         raise AssertionError("Invalid plan type")
 
@@ -719,7 +714,6 @@ def do_change_realm_plan_type(
             "_max_invites",
             "enable_spectator_access",
             "message_visibility_limit",
-            "upload_quota_gb",
         ]
     )
 

zerver/migrations/0507_rework_realm_upload_quota_gb.py

@@ -0,0 +1,21 @@
+# Generated by Django 4.2.10 on 2024-03-10 04:02
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("zerver", "0506_realm_require_unique_names"),
+    ]
+
+    operations = [
+        migrations.RemoveField(
+            model_name="realm",
+            name="upload_quota_gb",
+        ),
+        migrations.AddField(
+            model_name="realm",
+            name="custom_upload_quota_gb",
+            field=models.IntegerField(null=True),
+        ),
+    ]

zerver/models/realms.py

@@ -510,10 +510,9 @@ class Realm(models.Model):  # type: ignore[django-manager-missing] # django-stub
         BOT_CREATION_ADMINS_ONLY,
     ]
 
-    # See upload_quota_bytes; don't interpret upload_quota_gb directly.
     UPLOAD_QUOTA_LIMITED = 5
     UPLOAD_QUOTA_STANDARD = 50
-    upload_quota_gb = models.IntegerField(null=True)
+    custom_upload_quota_gb = models.IntegerField(null=True)
 
     VIDEO_CHAT_PROVIDERS = {
         "disabled": {
@@ -830,6 +829,27 @@ class Realm(models.Model):  # type: ignore[django-manager-missing] # django-stub
     def max_invites(self, value: Optional[int]) -> None:
         self._max_invites = value
 
+    @property
+    def upload_quota_gb(self) -> Optional[int]:
+        # See upload_quota_bytes; don't interpret upload_quota_gb directly.
+
+        if self.custom_upload_quota_gb is not None:
+            return self.custom_upload_quota_gb
+
+        plan_type = self.plan_type
+        if plan_type == Realm.PLAN_TYPE_PLUS:
+            return Realm.UPLOAD_QUOTA_STANDARD
+        elif plan_type == Realm.PLAN_TYPE_STANDARD:
+            return Realm.UPLOAD_QUOTA_STANDARD
+        elif plan_type == Realm.PLAN_TYPE_SELF_HOSTED:
+            return None
+        elif plan_type == Realm.PLAN_TYPE_STANDARD_FREE:
+            return Realm.UPLOAD_QUOTA_STANDARD
+        elif plan_type == Realm.PLAN_TYPE_LIMITED:
+            return Realm.UPLOAD_QUOTA_LIMITED
+        else:
+            raise AssertionError("Invalid plan type")
+
     def upload_quota_bytes(self) -> Optional[int]:
         if self.upload_quota_gb is None:
             return None

zerver/tests/test_realm.py

@@ -977,6 +977,17 @@ class RealmTest(ZulipTestCase):
         self.assertEqual(realm.upload_quota_gb, Realm.UPLOAD_QUOTA_STANDARD)
         self.assertEqual(realm.can_access_all_users_group_id, everyone_system_group.id)
 
+        # Test that custom_upload_quota_gb overrides the default upload_quota_gb
+        # implied by a plan and makes .upload_quota_gb be unaffacted by plan changes.
+        realm.custom_upload_quota_gb = 100
+        realm.save(update_fields=["custom_upload_quota_gb"])
+        do_change_realm_plan_type(realm, Realm.PLAN_TYPE_PLUS, acting_user=iago)
+        self.assertEqual(realm.plan_type, Realm.PLAN_TYPE_PLUS)
+        self.assertEqual(realm.upload_quota_gb, 100)
+
+        realm.custom_upload_quota_gb = None
+        realm.save(update_fields=["custom_upload_quota_gb"])
+
         do_change_realm_plan_type(realm, Realm.PLAN_TYPE_SELF_HOSTED, acting_user=iago)
         self.assertEqual(realm.plan_type, Realm.PLAN_TYPE_SELF_HOSTED)
         self.assertEqual(realm.max_invites, settings.INVITES_DEFAULT_REALM_DAILY_MAX)

zerver/tests/test_upload.py

@@ -604,8 +604,8 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase):
         d1_attachment = Attachment.objects.get(path_id=d1_path_id)
 
         realm = get_realm("zulip")
-        realm.upload_quota_gb = 1
-        realm.save(update_fields=["upload_quota_gb"])
+        realm.custom_upload_quota_gb = 1
+        realm.save(update_fields=["custom_upload_quota_gb"])
 
         # The size of StringIO("zulip!") is 6 bytes. Setting the size of
         # d1_attachment to realm.upload_quota_bytes() - 11 should allow
@@ -625,8 +625,8 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase):
         result = self.client_post("/json/user_uploads", {"file": d3})
         self.assert_json_error(result, "Upload would exceed your organization's upload quota.")
 
-        realm.upload_quota_gb = None
-        realm.save(update_fields=["upload_quota_gb"])
+        realm.custom_upload_quota_gb = None
+        realm.save(update_fields=["custom_upload_quota_gb"])
         result = self.client_post("/json/user_uploads", {"file": d3})
         self.assert_json_success(result)