diff --git a/apps/rbac/migrations/0004_auto_20211201_1901.py b/apps/rbac/migrations/0004_auto_20211201_1901.py index eafe37539..9d59d99fc 100644 --- a/apps/rbac/migrations/0004_auto_20211201_1901.py +++ b/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)] - 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) 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)] - for i, group in enumerate(grouped_members): + count = 0 + 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): diff --git a/utils/test_run_migrations.py b/utils/test_run_migrations.py index c45c46ec1..33c7c3c5c 100644 --- a/utils/test_run_migrations.py +++ b/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()