perf: 继续优化一波

apps/rbac/migrations/0004_auto_20211201_1901.py

@@ -1,3 +1,5 @@
+# Generated by Django 3.1.13 on 2021-12-01 11:01
+
 import time
 from django.db import migrations
 
@@ -8,35 +10,48 @@ def migrate_system_role_binding(apps, schema_editor):
     db_alias = schema_editor.connection.alias
     user_model = apps.get_model('users', 'User')
     role_binding_model = apps.get_model('rbac', 'SystemRoleBinding')
-    users = user_model.objects.using(db_alias).all()
 
-    grouped_users = [users[i:i+1000] for i in range(0, len(users), 1000)]
+    count = 0
-    for i, group in enumerate(grouped_users):
+    bulk_size = 1000
+    while True:
+        users = user_model.objects.using(db_alias) \
+            .only('role', 'id') \
+            .all()[count:count+bulk_size]
+        if not users:
+            break
+
         role_bindings = []
         start = time.time()
-        for user in group:
+        for user in users:
             role = BuiltinRole.get_system_role_by_old_name(user.role)
             role_binding = role_binding_model(scope='system', user_id=user.id, role_id=role.id)
             role_bindings.append(role_binding)
+
         role_binding_model.objects.bulk_create(role_bindings, ignore_conflicts=True)
         print("Create role binding: {}-{} using: {:.2f}s".format(
-            i*1000, i*1000 + len(group), time.time()-start
+            count, count + len(users), time.time()-start
         ))
+        count += len(users)
 
 
 def migrate_org_role_binding(apps, schema_editor):
     db_alias = schema_editor.connection.alias
     org_member_model = apps.get_model('orgs', 'OrganizationMember')
     role_binding_model = apps.get_model('rbac', 'RoleBinding')
-    members = org_member_model.objects.using(db_alias)\
-        .only('role', 'user_id', 'org_id')\
-        .all()
 
-    grouped_members = [members[i:i+1000] for i in range(0, len(members), 1000)]
+    count = 0
-    for i, group in enumerate(grouped_members):
+    bulk_size = 1000
+
+    while True:
+        members = org_member_model.objects.using(db_alias)\
+            .only('role', 'user_id', 'org_id')\
+            .all()[count:count+bulk_size]
+        if not members:
+            break
         role_bindings = []
         start = time.time()
-        for member in group:
+
+        for member in members:
             role = BuiltinRole.get_org_role_by_old_name(member.role)
             role_binding = role_binding_model(
                 scope='org',
@@ -47,8 +62,9 @@ def migrate_org_role_binding(apps, schema_editor):
             role_bindings.append(role_binding)
         role_binding_model.objects.bulk_create(role_bindings, ignore_conflicts=True)
         print("Create role binding: {}-{} using: {:.2f}s".format(
-            i*1000, i*1000 + len(group), time.time()-start
+            count, count + len(members), time.time()-start
         ))
+        count += len(members)
 
 
 class Migration(migrations.Migration):

utils/test_run_migrations.py

@@ -20,24 +20,32 @@ from django.db import migrations
 from rbac.builtin import BuiltinRole
 
 
-def migrate_some_binding(apps, schema_editor):
+def migrate_system_role_binding(apps, schema_editor):
     db_alias = schema_editor.connection.alias
     user_model = apps.get_model('users', 'User')
     role_binding_model = apps.get_model('rbac', 'SystemRoleBinding')
-    users = user_model.objects.using(db_alias).all()
 
-    grouped_users = [users[i:i+1000] for i in range(0, len(users), 1000)]
+    count = 0
-    for i, group in enumerate(grouped_users):
+    bulk_size = 1000
+    while True:
+        users = user_model.objects.using(db_alias) \
+                    .only('role', 'id') \
+                    .all()[count:count+bulk_size]
+        if not users:
+            break
+
         role_bindings = []
         start = time.time()
-        for user in group:
+        for user in users:
             role = BuiltinRole.get_system_role_by_old_name(user.role)
             role_binding = role_binding_model(scope='system', user_id=user.id, role_id=role.id)
             role_bindings.append(role_binding)
+
         role_binding_model.objects.bulk_create(role_bindings, ignore_conflicts=True)
         print("Create role binding: {}-{} using: {:.2f}s".format(
-            i*1000, i*1000 + len(group), time.time()-start
+            count, count + len(users), time.time()-start
         ))
+        count += len(users)
 
 
 class Migration(migrations.Migration):
@@ -47,14 +55,14 @@ class Migration(migrations.Migration):
     ]
 
     operations = [
-        migrations.RunPython(migrate_some_binding),
+        migrations.RunPython(migrate_system_role_binding),
     ]
 
 
 # ================== 添加到下方 ======================
 def main():
     schema_editor = connection.schema_editor()
-    migrate_some_binding(apps, schema_editor)
+    migrate_system_role_binding(apps, schema_editor)
 
 
 # main()