From 0464b1a9e69c2ca407b213a797652460b2620159 Mon Sep 17 00:00:00 2001 From: ibuler Date: Mon, 13 Jun 2022 14:22:07 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BC=98=E5=8C=96=E8=BF=81=E7=A7=BB=20?= =?UTF-8?q?rbac=20=E9=80=9F=E5=BA=A6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit perf: migrate --- apps/rbac/builtin.py | 30 ++++++---- .../migrations/0004_auto_20211201_1901.py | 55 ++++++++++------- utils/test_run_migrations.py | 60 +++++++++++++++++++ 3 files changed, 111 insertions(+), 34 deletions(-) create mode 100644 utils/test_run_migrations.py diff --git a/apps/rbac/builtin.py b/apps/rbac/builtin.py index b736ba226..34733dc59 100644 --- a/apps/rbac/builtin.py +++ b/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): diff --git a/apps/rbac/migrations/0004_auto_20211201_1901.py b/apps/rbac/migrations/0004_auto_20211201_1901.py index 36736f254..eafe37539 100644 --- a/apps/rbac/migrations/0004_auto_20211201_1901.py +++ b/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): diff --git a/utils/test_run_migrations.py b/utils/test_run_migrations.py new file mode 100644 index 000000000..c45c46ec1 --- /dev/null +++ b/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()