migrations: Fix unnecessary loop in 0602.

This commit is contained in:
Tim Abbott 2024-10-16 22:24:00 -07:00
parent 7a53db8498
commit c4f5cd6552
1 changed files with 14 additions and 23 deletions

View File

@ -3,7 +3,7 @@
from django.db import migrations, transaction from django.db import migrations, transaction
from django.db.backends.base.schema import BaseDatabaseSchemaEditor from django.db.backends.base.schema import BaseDatabaseSchemaEditor
from django.db.migrations.state import StateApps from django.db.migrations.state import StateApps
from django.db.models import Max, Min, OuterRef from django.db.models import Max, OuterRef
def remap_can_manage_all_groups_for_existing_realms( def remap_can_manage_all_groups_for_existing_realms(
@ -11,34 +11,25 @@ def remap_can_manage_all_groups_for_existing_realms(
) -> None: ) -> None:
Realm = apps.get_model("zerver", "Realm") Realm = apps.get_model("zerver", "Realm")
NamedUserGroup = apps.get_model("zerver", "NamedUserGroup") NamedUserGroup = apps.get_model("zerver", "NamedUserGroup")
BATCH_SIZE = 1000
max_id = NamedUserGroup.objects.aggregate(Max("id"))["id__max"] max_id = NamedUserGroup.objects.aggregate(Max("id"))["id__max"]
if max_id is None: if max_id is None:
# Do nothing if there are no user groups on the server. # Do nothing if there are no user groups on the server.
return return
lower_bound = NamedUserGroup.objects.aggregate(Min("id"))["id__min"] with transaction.atomic():
# Since can_manage_group, can_add_members_group, etc. have
while lower_bound <= max_id: # migrated to the nearest possible value from
upper_bound = lower_bound + BATCH_SIZE - 1 # user_group_edit_policy, we want to set
print(f"Processing batch {lower_bound} to {upper_bound} for NamedUserGroup") # can_manage_all_groups to the most restrictive setting
# previously possible. We've chosen administrators as the
with transaction.atomic(): # value here since the highest possible
# Since can_manage_group, can_add_members_group, etc. have # user_group_edit_policy was with role administrators.
# migrated to the nearest possible value from Realm.objects.update(
# user_group_edit_policy, we want to set can_manage_all_groups=NamedUserGroup.objects.filter(
# can_manage_all_groups to the most restrictive setting name="role:administrators", realm=OuterRef("id"), is_system_group=True
# previously possible. We've chosen administrators as the ).values("pk")
# value here since the highest possible )
# user_group_edit_policy was with role administrators.
Realm.objects.update(
can_manage_all_groups=NamedUserGroup.objects.filter(
name="role:administrators", realm=OuterRef("id"), is_system_group=True
).values("pk")
)
lower_bound += BATCH_SIZE
class Migration(migrations.Migration): class Migration(migrations.Migration):