perf: 继续优化一波

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)]
-    for i, group in enumerate(grouped_users):
+    count = 0
+    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()