perf: 优化迁移 rbac 速度

perf: migrate

apps/rbac/builtin.py

@@ -126,6 +126,8 @@ class BuiltinRole:
     org_user = PredefineRole(
         '7', ugettext_noop('OrgUser'), Scope.org, user_perms
     )
+    system_role_mapper = None
+    org_role_mapper = None
 
     @classmethod
     def get_roles(cls):
@@ -138,22 +140,24 @@ class BuiltinRole:
 
     @classmethod
     def get_system_role_by_old_name(cls, name):
-        mapper = {
-            'App': cls.system_component,
-            'Admin': cls.system_admin,
-            'User': cls.system_user,
-            'Auditor': cls.system_auditor
-        }
-        return mapper[name].get_role()
+        if not cls.system_role_mapper:
+            cls.system_role_mapper = {
+                'App': cls.system_component.get_role(),
+                'Admin': cls.system_admin.get_role(),
+                'User': cls.system_user.get_role(),
+                'Auditor': cls.system_auditor.get_role()
+            }
+        return cls.system_role_mapper[name]
 
     @classmethod
     def get_org_role_by_old_name(cls, name):
-        mapper = {
-            'Admin': cls.org_admin,
-            'User': cls.org_user,
-            'Auditor': cls.org_auditor,
-        }
-        return mapper[name].get_role()
+        if not cls.org_role_mapper:
+            cls.org_role_mapper = {
+                'Admin': cls.org_admin.get_role(),
+                'User': cls.org_user.get_role(),
+                'Auditor': cls.org_auditor.get_role(),
+            }
+        return cls.org_role_mapper[name]
 
     @classmethod
     def sync_to_db(cls, show_msg=False):

apps/rbac/migrations/0004_auto_20211201_1901.py

@@ -1,5 +1,4 @@
-# Generated by Django 3.1.13 on 2021-12-01 11:01
-
+import time
 from django.db import migrations
 
 from rbac.builtin import BuiltinRole
@@ -11,31 +10,45 @@ def migrate_system_role_binding(apps, schema_editor):
     role_binding_model = apps.get_model('rbac', 'SystemRoleBinding')
     users = user_model.objects.using(db_alias).all()
 
-    role_bindings = []
-    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)
+    grouped_users = [users[i:i+1000] for i in range(0, len(users), 1000)]
+    for i, group in enumerate(grouped_users):
+        role_bindings = []
+        start = time.time()
+        for user in group:
+            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
+        ))
 
 
 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).all()
-
-    role_bindings = []
-    for member in members:
-        role = BuiltinRole.get_org_role_by_old_name(member.role)
-        role_binding = role_binding_model(
-            scope='org',
-            user_id=member.user.id,
-            role_id=role.id,
-            org_id=member.org.id
-        )
-        role_bindings.append(role_binding)
-    role_binding_model.objects.bulk_create(role_bindings)
+    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)]
+    for i, group in enumerate(grouped_members):
+        role_bindings = []
+        start = time.time()
+        for member in group:
+            role = BuiltinRole.get_org_role_by_old_name(member.role)
+            role_binding = role_binding_model(
+                scope='org',
+                user_id=member.user_id,
+                role_id=role.id,
+                org_id=member.org_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
+        ))
 
 
 class Migration(migrations.Migration):

utils/test_run_migrations.py

@@ -0,0 +1,60 @@
+# Generated by Django 3.1.13 on 2021-12-01 11:01
+import os
+import sys
+import django
+import time
+
+app_path = '***** Change me *******'
+sys.path.insert(0, app_path)
+os.environ.setdefault("DJANGO_SETTINGS_MODULE", "jumpserver.settings")
+django.setup()
+
+from django.apps import apps
+from django.db import connection
+
+# ========================== 添加到需要测试的 migrations 上方 ==========================
+
+
+from django.db import migrations
+
+from rbac.builtin import BuiltinRole
+
+
+def migrate_some_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):
+        role_bindings = []
+        start = time.time()
+        for user in group:
+            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
+        ))
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('rbac', '0003_auto_20211130_1037'),
+    ]
+
+    operations = [
+        migrations.RunPython(migrate_some_binding),
+    ]
+
+
+# ================== 添加到下方 ======================
+def main():
+    schema_editor = connection.schema_editor()
+    migrate_some_binding(apps, schema_editor)
+
+
+# main()